def fuse_permute_bmm_and_gemm(
    sorted_graph: List[Tensor], workdir: str = None
) -> List[Tensor]:
    """Fuse [permute021 + bmm] and [permute(0, 1) + gemm].

    Note that for the latter fusion, we require that this pass takes
    place before any gemm + elementwise fusions.

    Parameters
    ----------
    sorted_graph : List[Tensor]
        Input graph
    workdir : str, optional
        working dir, by default None

    Returns
    -------
    List[Tensor]
        Fused graph
    """

    def _need_broadcast_gemm(op: Operator):
        if not op._attrs["op"].startswith("gemm"):
            return False
        inputs = op._attrs["inputs"]
        # cutlass's bmm assigns the batch size to grid_z, which
        # cannot exceeds 65535
        MAX_B_DIM_VAL = 65535

        def _valid_shape(shape: List[Union[IntImm, IntVar]]):
            b_dim = shape[0]
            if isinstance(b_dim, IntImm):
                b_dim_val = b_dim.value()
            else:
                b_dim_val = b_dim.upper_bound()
            return b_dim_val <= MAX_B_DIM_VAL

        input_shape_0 = inputs[0].shape()
        input_shape_1 = inputs[1].shape()
        if len(input_shape_0) != 3 and len(input_shape_1) != 3:
            return False
        if len(input_shape_0) == 3 and not _valid_shape(input_shape_0):
            return False
        if len(input_shape_1) == 3 and not _valid_shape(input_shape_1):
            return False
        return True

    def _is_transpose(op: Operator):
        if op._attrs["op"] != "permute":
            return False
        dims = op._attrs["dims"]
        return dims == [1, 0]

    permute_mm_patterns = (
        ([permute021, bmm_ccr], [bmm_rcr, bmm_crr], None, None),
        ([permute021, bmm_crr], [bmm_rrr, bmm_ccr], None, None),
        ([permute021, bmm_rcr], [bmm_ccr, bmm_rrr], None, None),
        ([permute021, bmm_rrr], [bmm_crr, bmm_rcr], None, None),
        ([permute021, gemm_rcr], [bmm_ccr, bmm_rrr], _need_broadcast_gemm, None),
        ([permute021, gemm_rrr], [bmm_crr, bmm_rcr], _need_broadcast_gemm, None),
        (
            [permute021, gemm_rcr_bias],
            [ops.gemm_universal.bmm_ccr_add, ops.gemm_universal.bmm_rrr_add],
            _need_broadcast_gemm,
            None,
        ),
        (
            [permute021, gemm_rrr_bias],
            [ops.gemm_universal.bmm_crr_add, None],
            _need_broadcast_gemm,
            None,
        ),
        ([permute, gemm_rcr], [None, gemm_rrr], None, _is_transpose),
        ([permute, gemm_rrr], [None, gemm_rcr], None, _is_transpose),
    )

    graph_transformed = True
    while graph_transformed:
        graph_transformed = False
        for source, targets, gemm_condition, permute_condition in permute_mm_patterns:
            fused, sorted_graph = _fuse_permute_impl(
                sorted_graph, source, targets, gemm_condition, permute_condition
            )
            graph_transformed |= fused

    return sorted_graph
Fuse [permute021 + bmm] and [permute(0, 1) + gemm].

    Note that for the latter fusion, we require that this pass takes
    place before any gemm + elementwise fusions.

    Parameters
    ----------
    sorted_graph : List[Tensor]
        Input graph
    workdir : str, optional
        working dir, by default None

    Returns
    -------
    List[Tensor]
        Fused graph
    
python/aitemplate/compiler/transform/fuse_permute_bmm_and_gemm.py