diff --git a/onnxruntime/python/tools/transformers/fusion_gelu.py b/onnxruntime/python/tools/transformers/fusion_gelu.py index 8626ca1482..6be5140c07 100644 --- a/onnxruntime/python/tools/transformers/fusion_gelu.py +++ b/onnxruntime/python/tools/transformers/fusion_gelu.py @@ -98,10 +98,13 @@ class FusionGelu(Fusion): return self.nodes_to_remove.extend(subgraph_nodes) - fused_node = helper.make_node("Gelu", inputs=[subgraph_input], outputs=[subgraph_output]) + fused_node = helper.make_node( + "Gelu", inputs=[subgraph_input], outputs=[subgraph_output], name=self.model.create_node_name("Gelu") + ) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) self.node_name_to_graph_name[fused_node.name] = self.this_graph_name + self.increase_counter("Gelu") return True def fuse_2(self, erf_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]: @@ -172,10 +175,13 @@ class FusionGelu(Fusion): return self.nodes_to_remove.extend(subgraph_nodes) - fused_node = helper.make_node("Gelu", inputs=[root_node.output[0]], outputs=[mul.output[0]]) + fused_node = helper.make_node( + "Gelu", inputs=[root_node.output[0]], outputs=[mul.output[0]], name=self.model.create_node_name("Gelu") + ) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) self.node_name_to_graph_name[fused_node.name] = self.this_graph_name + self.increase_counter("Gelu") return True def fuse_3(self, erf_node, input_name_to_nodes: Dict, output_name_to_node: Dict) -> Optional[bool]: @@ -243,8 +249,11 @@ class FusionGelu(Fusion): return self.nodes_to_remove.extend(subgraph_nodes) - fused_node = helper.make_node("Gelu", inputs=[root_node.output[0]], outputs=[last_mul.output[0]]) + fused_node = helper.make_node( + "Gelu", inputs=[root_node.output[0]], outputs=[last_mul.output[0]], name=self.model.create_node_name("Gelu") + ) fused_node.domain = "com.microsoft" self.nodes_to_add.append(fused_node) self.node_name_to_graph_name[fused_node.name] = self.this_graph_name + self.increase_counter("Gelu") return True diff --git a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py index 1ec5edf686..a10b61fdc3 100644 --- a/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py +++ b/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py @@ -24,14 +24,15 @@ class FusionSkipLayerNormalization(Fusion): model: OnnxModel, fused_op_type: str = "SkipLayerNormalization", search_op_types: str = "LayerNormalization", + shape_infer: bool = True, ): super().__init__(model, fused_op_type, search_op_types) - # Update shape inference is needed since other fusions might add new edge which does not have shape info yet. - self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7}, update=True) - - if self.shape_infer_helper is None: - # TODO(tianleiwu): support subgraph in shape inference or add broadcasting in SkipLayerNormalization op. - logger.warning("symbolic shape inference disabled or failed.") + if shape_infer: + # Update shape inference is needed since other fusions might add new edge which does not have shape info yet. + self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7}, update=True) + if self.shape_infer_helper is None: + # TODO(tianleiwu): support subgraph in shape inference or add broadcasting in SkipLayerNormalization op. + logger.warning("symbolic shape inference disabled or failed.") def fuse(self, node, input_name_to_nodes, output_name_to_node): add = self.model.get_parent(node, 0, output_name_to_node) @@ -56,18 +57,19 @@ class FusionSkipLayerNormalization(Fusion): # Root Mean Square Layer Normalization simplified = node.op_type == "SimplifiedLayerNormalization" - if self.shape_infer_helper is not None: - # TODO(tianleiwu): support broadcasting Skip shape (1, sequence_length, hidden_size) or (sequence_length, hidden_size) - if not self.shape_infer_helper.compare_shape(add.input[0], add.input[1]): - logger.debug( - "skip SkipLayerNormalization fusion since shape of inputs (%s, %s) are not same", - add.input[0], - add.input[1], - ) + if hasattr(self, "shape_infer_helper"): + if self.shape_infer_helper is not None: + # TODO(tianleiwu): support broadcasting Skip shape (1, sequence_length, hidden_size) or (sequence_length, hidden_size) + if not self.shape_infer_helper.compare_shape(add.input[0], add.input[1]): + logger.debug( + "skip SkipLayerNormalization fusion since shape of inputs (%s, %s) are not same", + add.input[0], + add.input[1], + ) + return + else: + logger.debug("skip SkipLayerNormalization fusion since symbolic shape inference failed") return - else: - logger.debug("skip SkipLayerNormalization fusion since symbolic shape inference failed") - return gather_path = self.model.match_parent_path(add, ["Gather"], [None]) if gather_path is not None and self.model.find_graph_input(gather_path[0].input[1]) is None: diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index a8fc6e6619..fe80a08829 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -63,9 +63,10 @@ class OnnxModel: return None - def input_name_to_nodes(self): + def input_name_to_nodes(self, exclude_subgraphs=False): input_name_to_nodes = {} - for node in self.nodes(): + nodes_to_search = self.nodes() if not exclude_subgraphs else self.model.graph.node + for node in nodes_to_search: for input_name in node.input: if input_name: # could be empty when it is optional if input_name not in input_name_to_nodes: @@ -74,9 +75,10 @@ class OnnxModel: input_name_to_nodes[input_name].append(node) return input_name_to_nodes - def output_name_to_node(self): + def output_name_to_node(self, exclude_subgraphs=False): output_name_to_node = {} - for node in self.nodes(): + nodes_to_search = self.nodes() if not exclude_subgraphs else self.model.graph.node + for node in nodes_to_search: for output_name in node.output: if output_name: # could be empty when it is optional output_name_to_node[output_name] = node @@ -906,6 +908,31 @@ class OnnxModel: if len(unused_nodes) > 0: logger.debug(f"Removed unused constant nodes: {len(unused_nodes)}") + def _get_subgraph_inputs_of_node(self, node): + """ + Get inputs to all nodes in all subgraphs of a node + """ + # Note: This function only handles one-level subgraphs of child nodes. + subgraph_nodes_inputs = set() + for attr in node.attribute: + if attr.type == AttributeProto.GRAPH: + child_nodes = attr.g.node + for child_node in child_nodes: + subgraph_nodes_inputs.update(child_node.input) + return subgraph_nodes_inputs + + def _get_subgraph_nodes_and_inputs(self, ops_with_graph_attrs): + """ + Get input names to all nodes in all subgraphs where subgraphs are + graph attributes of a node in the main graph + """ + subgraph_nodes = list(filter(lambda node: node.op_type in ops_with_graph_attrs, self.model.graph.node)) + subgraph_nodes_inputs = set() + for parent_node in subgraph_nodes: + subgraph_inputs_of_parent_node = self._get_subgraph_inputs_of_node(parent_node) + subgraph_nodes_inputs.update(subgraph_inputs_of_parent_node) + return subgraph_nodes, subgraph_nodes_inputs + def prune_graph(self, outputs=None, allow_remove_graph_inputs=True): """ Prune graph to keep only required outputs. It removes unnecessary nodes that are not linked @@ -918,13 +945,9 @@ class OnnxModel: allow_remove_graph_inputs (bool): allow remove graph inputs. """ - if len(self.graphs()) > 1: - # TODO(tianleiwu): handle subgraph - logger.debug("Skip prune_graph since graph has subgraph") - return - keep_outputs = [output.name for output in self.model.graph.output] if outputs is None else outputs + input_name_to_nodes_for_main_graph = self.input_name_to_nodes(exclude_subgraphs=True) output_name_to_node = self.output_name_to_node() def get_first_output(node): @@ -932,6 +955,29 @@ class OnnxModel: return node.output[0] return next(iter([o for o in node.output if o]), None) + if len(self.graphs()) > 1: + # Get input names for all nodes in all subgraphs + subgraph_nodes, subgraph_nodes_inputs = self._get_subgraph_nodes_and_inputs( + ops_with_graph_attrs={"Loop", "Scan", "If"} + ) + if len(subgraph_nodes) == 0: + # TODO: support other ops such as `BeamSearch` that have subgraphs as op attributes + logger.debug("Skip prune_graph since graph has subgraph") + return + + # For graphs with subgraphs, add dangling outputs from parent graph nodes to list of outputs to keep + for node in self.model.graph.node: + # TODO: This for-loop logic currently assumes that Loop/Scan/If nodes will not be + # pruned because their subgraphs are needed for computations. This might not be + # true in all cases. + if node in subgraph_nodes: + continue + + # Check if node output is an input of a subgraph node and not an input to a node in the main graph + for output in node.output: + if output in subgraph_nodes_inputs and output not in input_name_to_nodes_for_main_graph: + keep_outputs += [output] + # Keep track of nodes to keep. The key is first output of node, and the value is the node. output_to_node = {} @@ -956,7 +1002,7 @@ class OnnxModel: first_output = get_first_output(node) kept_node = output_to_node.get(first_output) - # Need double check the node since fused node might reuse output name of some nodes to be removed. + # Need to double check the node since fused node might reuse output name of some nodes to be removed. # It is slow to compare whole node, so we compare op_type first to avoid comparing node in most cases. if kept_node and kept_node.op_type == node.op_type and kept_node == node: nodes_to_keep.append(node) @@ -997,16 +1043,15 @@ class OnnxModel: def update_graph(self, verbose=False, allow_remove_graph_inputs=False): graph = self.model.graph - remaining_input_names = [] + remaining_input_names = set() for node in graph.node: if node.op_type in ["Loop", "Scan", "If"]: - # TODO: handle inner graph - logger.debug(f"Skip update_graph since graph has operator: {node.op_type}") - return + # Add input names of nodes in subgraphs + subgraph_inputs_of_node = self._get_subgraph_inputs_of_node(node) + remaining_input_names.update(subgraph_inputs_of_node) + if node.op_type != "Constant": - for input_name in node.input: - if input_name not in remaining_input_names: - remaining_input_names.append(input_name) + remaining_input_names.update(node.input) if verbose: logger.debug(f"remaining input names: {remaining_input_names}") diff --git a/onnxruntime/python/tools/transformers/onnx_model_bert.py b/onnxruntime/python/tools/transformers/onnx_model_bert.py index ad51c1cce0..26464fc328 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_bert.py +++ b/onnxruntime/python/tools/transformers/onnx_model_bert.py @@ -115,8 +115,8 @@ class BertOnnxModel(OnnxModel): fusion = FusionSimplifiedLayerNormalization(self) fusion.apply() - def fuse_skip_layer_norm(self): - fusion = FusionSkipLayerNormalization(self) + def fuse_skip_layer_norm(self, shape_infer=True): + fusion = FusionSkipLayerNormalization(self, shape_infer=shape_infer) fusion.apply() def fuse_skip_simplified_layer_norm(self): @@ -344,7 +344,7 @@ class BertOnnxModel(OnnxModel): self.fuse_reshape() if (options is None) or options.enable_skip_layer_norm: - self.fuse_skip_layer_norm() + self.fuse_skip_layer_norm(options.enable_shape_inference) self.fuse_skip_simplified_layer_norm() if (options is None) or options.enable_rotary_embeddings: diff --git a/onnxruntime/python/tools/transformers/onnx_model_clip.py b/onnxruntime/python/tools/transformers/onnx_model_clip.py index 32bddc3ca1..388d058c78 100644 --- a/onnxruntime/python/tools/transformers/onnx_model_clip.py +++ b/onnxruntime/python/tools/transformers/onnx_model_clip.py @@ -24,6 +24,7 @@ class ClipOnnxModel(BertOnnxModel): op_count = {} ops = [ "Attention", + "Gelu", "LayerNormalization", "QuickGelu", "SkipLayerNormalization",