def _fuse_gemm_rcr_bias_swish(sorted_graph: List[Tensor]) -> List[Tensor]:
    """
    gemm_rcr_bias_swish(A, B) is equivalent to:
        x = gemm_rcr_bias(A, B)
        x1 = sigmoid(x)
        return elementwise(MUL)(x, x1)
    """
    new_sorted_graph = []

    to_remove = set()
    for tensor in sorted_graph:
        if tensor in to_remove:
            continue
        new_sorted_graph.append(tensor)

        if tensor._attrs["is_output"]:
            continue

        gemm_op = extract_only_one_op(tensor._attrs["src_ops"])
        if gemm_op is None:
            continue
        if gemm_op._attrs["op"] != "gemm_rcr_bias":
            continue

        dst_op = list(tensor._attrs["dst_ops"])
        if len(dst_op) != 2:
            continue
        swish_tensor = None
        for idx in range(2):
            other_idx = (idx + 1) % 2
            if is_elementwise_type(dst_op[idx], FuncEnum.SIGMOID):
                if not is_elementwise_type(dst_op[other_idx], FuncEnum.MUL):
                    continue

                is_swish = False
                output = dst_op[idx]._attrs["outputs"][0]
                mul_inputs = dst_op[other_idx]._attrs["inputs"]
                if mul_inputs[0] == output and mul_inputs[1] == tensor:
                    is_swish = True
                if mul_inputs[1] == output and mul_inputs[0] == tensor:
                    is_swish = True
                if not is_swish:
                    continue

                swish_tensor = dst_op[other_idx]._attrs["outputs"][0]
                break

        if swish_tensor is None:
            continue

        gemm_inputs = gemm_op._attrs["inputs"]
        remove_dst_op_from_tensor(gemm_inputs, gemm_op)
        # Output of sigmoid and final mul of swish.
        to_remove.add(dst_op[0]._attrs["outputs"][0])
        to_remove.add(dst_op[1]._attrs["outputs"][0])

        new_tensor = gemm_rcr_bias_swish()(*gemm_inputs)
        copy_tensor_attributes(new_tensor, swish_tensor)
        replace_tensor(swish_tensor, new_tensor)
        new_sorted_graph[-1] = new_tensor

    return sanitize_sorted_graph(new_sorted_graph)

    gemm_rcr_bias_swish(A, B) is equivalent to:
        x = gemm_rcr_bias(A, B)
        x1 = sigmoid(x)
        return elementwise(MUL)(x, x1)
    
python/aitemplate/compiler/transform/fuse_mm_elementwise.py