Class RewritePatternManager

tensorrt_llm/graph_rewriting.py:247–274

class RewritePatternManager(_PatternManager):
    def rewrite(self, net: Network, args=None):
        modified = True
        # TODO: we can optimize this by asking TRT to expose a graph iterator consistent even after
        # the graph is modified.
        while modified:
            modified = False
            # Since the graph iterator is hold by the underlying INetwork, we can only rebuild the
            # graph cache and match the nodes again.
            for layer in net.get_layers():
                if layer.is_removed():
                    continue
                for profit, pattern in sorted(self.patterns.values(), key=lambda x: x[0]):
                    pattern.args = args

                    if pattern.root_layer is not None and layer.type not in pattern.root_layer:
                        continue
                    if pattern._separate_match_rewrite:
                        if pattern.match(layer):
                            pattern.rewrite(layer)
                            modified = True
                    else:
                        if pattern.match_and_rewrite(layer):
                            modified = True

    @staticmethod
    def instance():
        return _global_rewrite_pattern_manager


class AnalysisPatternManager(_PatternManager):