def _transform_odd_alignment(
    sorted_graph: List[Tensor],
    permutable_pairs: Dict[str, Tuple[Operator, Operator, Operator]],
) -> List[Tensor]:
    """
    This function tries to insert new permute021() ops before gemms when applicable.

    For input tensors with odd alignments, either permutation or padding are needed
    to transform an odd alignment to an even alignment, so that SM80 cutlass kernels
    can be used.
    This function decides between permutation and padding by estimating cost of all
    the options, and selects the one with minimal cost. If adding permutation costs
    less, this function inserts new permute021 kernels before gemm ops. Otherwise
    this function does nothing and apply_padding pass kicks in.

    Cost of permutation is the total element of matrix.
    Cost of padding is computed with the following rules:
        1) If the connecting op cannot be fused with padding, the cost is the total
           elements of matrix.
        2) If the connecting op can be fused with padding, the cost would be the
           elements of the additional zeros padded.
    One special case is if the input is constant, both permutation and padding are
    free since we have constant folding pass.

    Parameters
    ----------
    sorted_graph : List[Tensor]
        AIT graph to run fusion
    permutable_pairs: Dict[str, Tuple[Operator, Operator, Operator]]
        A dictionary where key is the op that might accept permuted inputs,
        and value is a tuple of len-3 which correspond to ops (permute first input,
        permute second input, permute both first and second inputs)
    """

    new_sorted_graph = []
    permuted_inputs = {}
    for tensor in sorted_graph:
        new_sorted_graph.append(tensor)

        src_op = extract_only_one_op(tensor._attrs["src_ops"])
        if src_op is None:
            continue

        op_type = src_op._attrs["op"]
        if op_type not in permutable_pairs:
            continue
        # FIXME: This pass only works for half type. We may need to change it to
        # work with other types such as int8 later. Note that for float type, it
        # is safe to skip, because gemm/bmm with float inputs always meet alignment
        # requirements.
        if src_op._attrs["inputs"][0].dtype() != "float16":
            continue

        perm_type = ([False, False], [False, True], [True, False], [True, True])
        permute_input = [False, False]
        best_cost = inf
        for p in perm_type:
            perm_cost = _compute_required_flops(src_op, p[0], p[1])
            if perm_cost < best_cost:
                permute_input = p
                best_cost = perm_cost

        if not permute_input[0] and not permute_input[1]:
            continue

        inputs = src_op._attrs["inputs"]
        new_inputs = list(inputs)
        for idx in range(2):
            if permute_input[idx]:
                if inputs[idx] in permuted_inputs:
                    permuted_input = permuted_inputs[inputs[idx]]
                else:
                    input_shape = inputs[idx].shape()
                    if len(input_shape) == 2:
                        expanded_input = unsqueeze(0)(inputs[idx])
                        new_sorted_graph.insert(-1, expanded_input)
                        permuted_input = permute021()(expanded_input)
                    else:
                        permuted_input = permute021()(inputs[idx])
                    new_sorted_graph.insert(-1, permuted_input)
                    permuted_inputs[inputs[idx]] = permuted_input
                new_inputs[idx] = permuted_input

        if permute_input[0] and permute_input[1]:
            new_tensor = permutable_pairs[op_type][2]()(*new_inputs)
        elif permute_input[0]:
            new_tensor = permutable_pairs[op_type][0]()(*new_inputs)
        elif permute_input[1]:
            new_tensor = permutable_pairs[op_type][1]()(*new_inputs)
        copy_tensor_attributes(new_tensor, tensor)
        copy_src_op_attributes(new_tensor, tensor)
        replace_tensor(tensor, new_tensor)

        remove_dst_op_from_tensor(inputs, src_op)
        remove_tensor_from_sorted_graph(tensor)

        new_sorted_graph[-1] = new_tensor

    new_sorted_graph = toposort(new_sorted_graph)
    return sanitize_sorted_graph(new_sorted_graph)

    This function tries to insert new permute021() ops before gemms when applicable.

    For input tensors with odd alignments, either permutation or padding are needed
    to transform an odd alignment to an even alignment, so that SM80 cutlass kernels
    can be used.
    This function decides between permutation and padding by estimating cost of all
    the options, and selects the one with minimal cost. If adding permutation costs
    less, this function inserts new permute021 kernels before gemm ops. Otherwise
    this function does nothing and apply_padding pass kicks in.

    Cost of permutation is the total element of matrix.
    Cost of padding is computed with the following rules:
        1) If the connecting op cannot be fused with padding, the cost is the total
           elements of matrix.
        2) If the connecting op can be fused with padding, the cost would be the
           elements of the additional zeros padded.
    One special case is if the input is constant, both permutation and padding are
    free since we have constant folding pass.

    Parameters
    ----------
    sorted_graph : List[Tensor]
        AIT graph to run fusion
    permutable_pairs: Dict[str, Tuple[Operator, Operator, Operator]]
        A dictionary where key is the op that might accept permuted inputs,
        and value is a tuple of len-3 which correspond to ops (permute first input,
        permute second input, permute both first and second inputs)
    
python/aitemplate/compiler/transform/transform_odd_alignment.py