def _merge_parallel_gemm_concat(
    gemm_ops: List[Operator], cat_op: Operator, sorted_graph: List[Tensor]
):
    """merge parallel gemm ops and the following concat op together"""
    # clear gemm_inputs dst_ops
    _clear_gemm_inputs_dst_ops(gemm_ops)

    inputs, weights, bias = _group_gemm_inputs(gemm_ops)

    n, k = weights[0].shape()[0].value(), weights[0].shape()[1].value()
    b = len(weights)

    dtype = inputs[0].dtype()
    rcr_align = default_align_ab(k, k, dtype)
    rrr_align = default_align_ab(k, n, dtype)

    use_rcr = rcr_align > rrr_align

    # create new subgraph
    bmm_input_cat = ops.concatenate()(inputs, dim=-1)
    bmm_input = ops.reshape()(bmm_input_cat, [-1, b, k])

    bmm_weight_cat = ops.concatenate()(weights, dim=0)
    bmm_weight_reshape = ops.reshape()(bmm_weight_cat, [b, n, k])
    bmm_weight = bmm_weight_reshape if use_rcr else ops.permute021()(bmm_weight_reshape)

    bmm_bias_cat = ops.concatenate()(bias, dim=0)
    bmm_bias = ops.reshape()(bmm_bias_cat, [b, n])

    if use_rcr:
        bmm = ops.perm102_bmm_rcr_bias()(bmm_input, bmm_weight, bmm_bias)
    else:
        bmm = ops.perm102_bmm_rrr_bias()(bmm_input, bmm_weight, bmm_bias)
    bmm_reshape = ops.reshape()(bmm, [-1, b * n])

    num_cat_inputs = len(cat_op._attrs["inputs"])
    cat_output = cat_op._attrs["outputs"][0]
    if len(gemm_ops) == num_cat_inputs:
        # fuse with concat op completely
        transform_utils.replace_tensor(cat_output, bmm_reshape)

        # if cat_output was the only output of the graph, we must
        # append the new graph output to the graph
        sorted_graph.append(bmm_reshape)

    else:
        # bmm_reshape now replaces num_cat_inputs cat inputs
        begin_idx = _get_gemm_output_idx_in_cat_inputs(gemm_ops[0], cat_op)
        end_idx = _get_gemm_output_idx_in_cat_inputs(gemm_ops[-1], cat_op)

        old_inputs = cat_op._attrs["inputs"]
        new_inputs = old_inputs[:begin_idx] + [bmm_reshape] + old_inputs[end_idx + 1 :]

        assert all(
            cat_op._attrs["input_masks"]
        ), "The input_pasts of cat_op must be all True"

        cat_op._attrs["inputs"] = new_inputs
        cat_op._attrs["input_accessors"] = [TensorAccessor(t) for t in new_inputs]
        cat_op._attrs["original_inputs"] = list(new_inputs)
        cat_op._attrs["input_masks"] = [True] * len(new_inputs)

        bmm_reshape._attrs["dst_ops"].add(cat_op)

    for gemm_op in gemm_ops:
        gemm_outputs = gemm_op._attrs["outputs"]
        assert len(gemm_outputs) == 1
        transform_utils.remove_tensor_from_sorted_graph(gemm_outputs[0])
merge parallel gemm ops and the following concat op together
python/aitemplate/compiler/transform/fuse_parallel_gemms.py