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

    into:
    # x: [m, b, k], w: [b, k, n], b: [b, n]
    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 constant folding:
    x = concatenate()([x1, x2], dim=-1).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

    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 = concatenate()([x1, x2], dim=-1).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.
    For the extra concat op to cat x1 and x2 together, we require that the ops
    that produce x1 and x2 write directly to the output of concat.
    It is required that all the gemm ops have the same problem sizes and 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.

    For odd k/n, we rely on apply_padding pass to add padding to X and W.
    The overall perf may be better or worse depending on problem sizes.

    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:
        # check cat op
        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_supported_op)

        for gemm_ops in all_groups:
            # TODO: 2 is arbitrarily chosen. More benchmarks with real models
            # are needed to find the best criteria. If gemms are not fused here,
            # they can be grouped into group gemms with other gemms.
            if len(gemm_ops) >= 2:
                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
# x1: [m, k], w1: [n, k], b1: [n]
y1 = gemm_rcr_bias()(x1, w1, b1)
y2 = gemm_rcr_bias()(x2, w2, b1)
y3 = concatenate()([x1, x2], dim=-1)

into:
# x: [m, b, k], w: [b, k, n], b: [b, n]
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 constant folding:
x = concatenate()([x1, x2], dim=-1).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

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 = concatenate()([x1, x2], dim=-1).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.
For the extra concat op to cat x1 and x2 together, we require that the ops
that produce x1 and x2 write directly to the output of concat.
It is required that all the gemm ops have the same problem sizes and 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.

For odd k/n, we rely on apply_padding pass to add padding to X and W.
The overall perf may be better or worse depending on problem sizes.

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