def fuse_single_source_parallel_gemms(
    sorted_graph: List[Tensor], workdir: str = None
) -> List[Tensor]:
    """This pass fuses patterns like
    # x: [m, k], w_i: [n_i, k], b_i: [n_i]
    y1 = gemm_rcr_bias()(x, w1, b1)
    y2 = gemm_rcr_bias()(x, w2, b2)
    ...

    into:
    # x: [m, k], w: [sum(n_i), k], b: [sum(n_i)]
    w = concatenate()([w1, w2], dim=0)
    b = concatenate()([b1, b2], dim=0)
    y = gemm_rcr_bias()(x, w, b)
    y1, y2 = split()(y, n)

    For w and b, we rely on constant folding to preprocess them.
    y1 and y2 would be written directly from y's op.
    It is required that all the gemm ops have the same layouts.

    On graph pass ordering, we need to make sure this pass runs before
    any other pass that modifies gemm and concat input/output TensorAccessors.

    Args:
        sorted_graph (List[Tensor]): a sorted list of tensors

    Returns:
        List[Tensor]: the transformed graph with all ops sorted
    """

    # Extract gemm_rcr_bias pattern first.
    sorted_graph = transform_simple_fusion_patterns(
        sorted_graph, get_gemm_rcr_bias_patterns()
    )

    applied = True
    while applied:
        applied, sorted_graph = _fuse_single_source_parallel_gemms(sorted_graph)

    return sorted_graph
This pass fuses patterns like
    # x: [m, k], w_i: [n_i, k], b_i: [n_i]
    y1 = gemm_rcr_bias()(x, w1, b1)
    y2 = gemm_rcr_bias()(x, w2, b2)
    ...

    into:
    # x: [m, k], w: [sum(n_i), k], b: [sum(n_i)]
    w = concatenate()([w1, w2], dim=0)
    b = concatenate()([b1, b2], dim=0)
    y = gemm_rcr_bias()(x, w, b)
    y1, y2 = split()(y, n)

    For w and b, we rely on constant folding to preprocess them.
    y1 and y2 would be written directly from y's op.
    It is required that all the gemm ops have the same layouts.

    On graph pass ordering, we need to make sure this pass runs before
    any other pass that modifies gemm and concat input/output TensorAccessors.

    Args:
        sorted_graph (List[Tensor]): a sorted list of tensors

    Returns:
        List[Tensor]: the transformed graph with all ops sorted
    
python/aitemplate/compiler/transform/fuse_parallel_gemms.py