def _group_ops_by_type(
    sorted_graph: List[Tensor], op_type: str, workdir: str = None
) -> bool:
    """Find and fuse all groups of ops that can be fused together.
    Each group is replaced with 1 group op.

    Args:
        sorted_graph (List[Tensor]): Topologically sorted input graph
        op_type (str): The type of op to be grouped

    Returns:
        True if we fused any group.

    The algorithm can be described as:
    0) Let groups = []
    1) Do dfs and get the dependency graph, `G: {op: all the descendants of op with op_type}`.
       G only contains ops with type == op_type. G does not need to be in topological order.
    2) Get all the candidate ops, `grouped: {op: flag}`, for group fusion. The flag
       denotes whether this op is grouped or not and is initialized to False. We need to
       filter out ops that are not eligible such as gemm ops with large m/n/k or odd
       alignment, or layernorm ops without gamma/beta etc. Later, we set all grouped ops
       in grouped to True. Ops that can't be grouped with any other ops are also set to True.
       **Ops in `grouped` must be in topological order.**
    3) For group gemm, we group all gemms following the same split op together so the
       split op can be eliminated. Due to the high cost of split, we don't apply the same
       m/n/k filter to these gemm ops. These ops are removed from `grouped`.
    4) Let op_set = set(grouped.keys()), all the ops available for grouping.
    5) For every op in grouped:
            If grouped[candidate] is True, continue to next op.
            Remove op from op_set. Because ops in grouped are topologically sorted, this
                guarantees that op_set won't contain any ancestors of op.
            Get op candidates that can be potentially grouped with op.
            candidates = op_set - {G[op]} where G[op] is descendant of op.
            Sort candidates by name (same as topological order).
            for every candidate in candidates:
                Check if op and candidate op can be grouped together.
                    - Check for dependency
                    - Check for op compatibility
                If yes:
                    Group them together and remove candidate from op_set.
                    Merge descendants of candidate op to descendants of op.
                    Set grouped[candidate] = True
            If the final group is >= 2, add them to `groups`
        Set grouped[op] = True
    """

    # TODO: as an optimization, we may keep using the same dependency_graph
    # through all the group passes and keep updating it
    dependency_graph = _get_dependency_graph(sorted_graph, op_type)

    # There is no op with op_type in the graph
    if len(dependency_graph) == 0:
        return False

    if workdir:
        _dump_dependency_graph(dependency_graph, op_type, "filtered", workdir)

    f_filter_op = _get_op_filter(op_type)
    f_check_ops_are_compatible = _get_op_checker(op_type)
    is_layernorm = op_type.startswith("layernorm")
    f_fuse_ops = _fuse_layernorm_ops if is_layernorm else _fuse_gemm_ops

    sorted_ops = graph_utils.get_sorted_ops(sorted_graph)

    # grouped: {key: op, value: whether this op has been grouped or not}
    # ops in grouped must be in topological order
    grouped = _get_sorted_candidate_ops(sorted_ops, op_type, f_filter_op)
    assert len(grouped) <= len(dependency_graph)

    groups = []

    # applies to group gemms only
    split_groups = _group_split_outputs_together(sorted_graph, sorted_ops, op_type)
    for group in split_groups:
        groups.append(group)
        for op in group:
            if op in grouped:
                del grouped[op]

    # the set of ops available for group fusion
    op_set = set(grouped.keys())

    for op, visited in grouped.items():
        if visited:
            op_set.discard(op)
            continue
        descendants = dependency_graph[op]
        op_set.discard(op)

        candidates = op_set - descendants

        def get_op_number(op: Operator) -> int:
            op_name = op._attrs["name"]
            last_idx = op_name.rfind("_")
            return int(op_name[last_idx + 1 :])

        # Sort by topological order. Sorting by names guarantees topological order
        # because of how the name_graph pass works
        group_candidates = sorted(candidates, key=lambda x: get_op_number(x))

        group = [op]
        for candidate in group_candidates:
            if candidate in descendants:
                continue

            if (
                _ops_have_same_num_inputs(op, candidate)
                and _check_op_num_outputs(op, 1)
                and f_check_ops_are_compatible(op, candidate)
            ):
                group.append(candidate)
                grouped[candidate] = True

                op_set.discard(candidate)

                # must merge descendants together
                descendants.update(dependency_graph[candidate])
        # remove any op that may introduce a cycle because of grouping ops
        group_op_idx = 0
        while group_op_idx < len(group):
            grouped_op = group[group_op_idx]
            if _has_cycle(grouped_op, group):
                del group[group_op_idx]
            else:
                group_op_idx += 1

        # We fuse each group right after we form it. Otherwise, _has_cycle may
        # miss cycles within groups. For example, see the graph below:
        #
        #        A --> C ---
        #                  |
        #    --> B --> D   |
        #    |             |
        #    --- X --> M   |
        #                  |
        #        Y --> N <--
        #
        # If we fuse (A, B) and (X, Y) at the same time, we would end up with a
        # cycle between the fused op (A, B) and (X, Y). On the other hand, if we
        # fuse (A, B) first, and then check _has_cycle before fusing (X, Y), we
        # will be able to detect the cycle.
        if len(group) > _MAX_LAYERNORM_GROUP and op_type.startswith("layernorm"):
            new_groups = _break_layernorm_groups(group)
            for new_group in new_groups:
                f_fuse_ops(new_group, sorted_graph)
            groups.extend(new_groups)
        elif len(group) >= 2:
            f_fuse_ops(group, sorted_graph)
            groups.append(group)

        grouped[op] = True

    if workdir:
        _dump_groups(groups, op_type, workdir)

    return len(groups) > 0
Find and fuse all groups of ops that can be fused together.
Each group is replaced with 1 group op.

Args:
sorted_graph (List[Tensor]): Topologically sorted input graph
op_type (str): The type of op to be grouped

Returns:
True if we fused any group.

The algorithm can be described as:
0) Let groups = []
1) Do dfs and get the dependency graph, `G: {op: all the descendants of op with op_type}`.
G only contains ops with type == op_type. G does not need to be in topological order.
2) Get all the candidate ops, `grouped: {op: flag}`, for group fusion. The flag
denotes whether this op is grouped or not and is initialized to False. We need to
filter out ops that are not eligible such as gemm ops with large m/n/k or odd
alignment, or layernorm ops without gamma/beta etc. Later, we set all grouped ops
in grouped to True. Ops that can't be grouped with any other ops are also set to True.
**Ops in `grouped` must be in topological order.**
3) For group gemm, we group all gemms following the same split op together so the
split op can be eliminated. Due to the high cost of split, we don't apply the same
m/n/k filter to these gemm ops. These ops are removed from `grouped`.
4) Let op_set = set(grouped.keys()), all the ops available for grouping.
5) For every op in grouped:
If grouped[candidate] is True, continue to next op.
Remove op from op_set. Because ops in grouped are topologically sorted, this
guarantees that op_set won't contain any ancestors of op.
Get op candidates that can be potentially grouped with op.
candidates = op_set - {G[op]} where G[op] is descendant of op.
Sort candidates by name (same as topological order).
for every candidate in candidates:
Check if op and candidate op can be grouped together.
- Check for dependency
- Check for op compatibility
If yes:
Group them together and remove candidate from op_set.
Merge descendants of candidate op to descendants of op.
Set grouped[candidate] = True
If the final group is >= 2, add them to `groups`
Set grouped[op] = True
python/aitemplate/compiler/transform/fuse_group_ops.py