def _fuse_split_parallel_gemm_concat(sorted_graph: List[Tensor]) -> List[Tensor]:
    """This pass fuses patterns like
    # x: [m, b * k], w1: [n, k], b1: [n]
    x1, x2 = split()(x, k, dim=-1)
    y1 = gemm_rcr_bias()(x1, w1, b1)
    y2 = gemm_rcr_bias()(x2, w2, b1)
    y3 = concatenate()([x1, x2], dim=-1)

    first into:
    # x: [m, b, k], w: [b, k, n], b: [b, n]
    x1, x2 = split()(x, k, dim=-1)
    x = concatenate()([x1, x2], dim=-1).reshape([-1, b, n])
    w = concatenate()([w1, w2], dim=-1).reshape([b, n, k]).permute([0, 2, 1])
    b =  concatenate()([b1, b2], dim=-1).reshape([b, n])
    y3 = perm102_bmm_rrr_bias()(x, w, b).reshape([-1, b * n])

    after transform_memory_ops pass:
    x = x.reshape([-1, b, n])
    w = concatenate()([w1, w2], dim=-1).reshape([b, n, k]).permute([0, 2, 1])
    b =  concatenate()([b1, b2], dim=-1).reshape([b, n])
    y3 = perm102_bmm_rrr_bias()(x, w, b).reshape([-1, b * n])

    after constant folding:
    x = x.reshape([-1, b, n])
    y3 = perm102_bmm_rrr_bias()(x, w, b).reshape([-1, b * n])

    If rcr layout has better alignment than rrr, it will tranform the graph into

    # x: [m, b, k], w: [b, n, k], b: [b, n]
    x = concatenate()([x1, x2], dim=-1).reshape([-1, b, n])
    y3 = perm102_bmm_rcr_bias()(x, w, b).reshape([-1, b * n])

    If there are other inputs to the concat op, such as

    x1, x2 = split()(x, k, dim=-1)
    y1 = gemm_rcr_bias()(x1, w1, b1)
    y2 = gemm_rcr_bias()(x2, w2, b1)
    y3 = concatenate()([y1, y2, x3, x4], dim=-1)

    The graph is transformed into
    x = x.reshape([-1, b, n])
    y3 = perm102_bmm_rrr_bias()(x, w, b).reshape([-1, b * n])
    y4 = concatenate()([y3, x3, x4], dim=-1)

    y3 will write into the y4 directly through concat fusion.

    For w and b, we rely on constant folding to preprocess them.
    It is required that all the gemm ops have the same problem sizes and layouts.
    We also check that all the gemm inputs come from the same split op.

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

    Returns:
        List[Tensor]: the transformed graph with all ops sorted
    """
    sorted_ops = graph_utils.get_sorted_ops(sorted_graph)
    fusion_groups = []
    for op in sorted_ops:
        if op._attrs["op"] != "concatenate":
            continue
        if not _check_cat_op(op):
            continue

        cat_inputs = op._attrs["inputs"]

        all_groups = _find_parallel_gemm_ops(cat_inputs, _is_split_op)

        for gemm_ops in all_groups:
            if not _from_same_src_op(gemm_ops, "split"):
                continue
            fusion_groups.append([gemm_ops, op])

    for gemm_ops, cat_op in fusion_groups:
        _merge_parallel_gemm_concat(gemm_ops, cat_op, sorted_graph)
    sorted_graph = toposort(sorted_graph)
    return transform_utils.sanitize_sorted_graph(sorted_graph)
This pass fuses patterns like
# x: [m, b * k], w1: [n, k], b1: [n]
x1, x2 = split()(x, k, dim=-1)
y1 = gemm_rcr_bias()(x1, w1, b1)
y2 = gemm_rcr_bias()(x2, w2, b1)
y3 = concatenate()([x1, x2], dim=-1)

first into:
# x: [m, b, k], w: [b, k, n], b: [b, n]
x1, x2 = split()(x, k, dim=-1)
x = concatenate()([x1, x2], dim=-1).reshape([-1, b, n])
w = concatenate()([w1, w2], dim=-1).reshape([b, n, k]).permute([0, 2, 1])
b =  concatenate()([b1, b2], dim=-1).reshape([b, n])
y3 = perm102_bmm_rrr_bias()(x, w, b).reshape([-1, b * n])

after transform_memory_ops pass:
x = x.reshape([-1, b, n])
w = concatenate()([w1, w2], dim=-1).reshape([b, n, k]).permute([0, 2, 1])
b =  concatenate()([b1, b2], dim=-1).reshape([b, n])
y3 = perm102_bmm_rrr_bias()(x, w, b).reshape([-1, b * n])

after constant folding:
x = x.reshape([-1, b, n])
y3 = perm102_bmm_rrr_bias()(x, w, b).reshape([-1, b * n])

If rcr layout has better alignment than rrr, it will tranform the graph into

# x: [m, b, k], w: [b, n, k], b: [b, n]
x = concatenate()([x1, x2], dim=-1).reshape([-1, b, n])
y3 = perm102_bmm_rcr_bias()(x, w, b).reshape([-1, b * n])

If there are other inputs to the concat op, such as

x1, x2 = split()(x, k, dim=-1)
y1 = gemm_rcr_bias()(x1, w1, b1)
y2 = gemm_rcr_bias()(x2, w2, b1)
y3 = concatenate()([y1, y2, x3, x4], dim=-1)

The graph is transformed into
x = x.reshape([-1, b, n])
y3 = perm102_bmm_rrr_bias()(x, w, b).reshape([-1, b * n])
y4 = concatenate()([y3, x3, x4], dim=-1)

y3 will write into the y4 directly through concat fusion.

For w and b, we rely on constant folding to preprocess them.
It is required that all the gemm ops have the same problem sizes and layouts.
We also check that all the gemm inputs come from the same split op.

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