def apply_padding(sorted_graph: List[Tensor], workdir: str = None) -> List[Tensor]:
    """
    Applies padding to gemms to use SM80 kernels.
    SM80 kernels require min_alignment == 2.
    """

    visited_ops = set()
    new_sorted_graph = []
    for tensor in sorted_graph:
        new_tensor_list = [tensor]
        src_ops = tensor.src_ops()
        for op in src_ops:
            if op in visited_ops:
                continue

            # Exclude special gemm kernels.
            if (
                not isinstance(op, gemm)
                or isinstance(op, ops.gemm_rrr_small_nk)
                or isinstance(op, ops.bmm_rcr_n1)
                or isinstance(op, ops.bmm_rrr_k1_tanh)
                or "permute" in op._attrs["op"]
            ):
                continue

            # This pass only works for gemm or bmm. group_gemm is not supported.
            # We don't need to padd our special kernel bmm_rcr_n1, which does
            # not have any alignment constraint.
            op_name = op._attrs["name"]
            if op_name.startswith(("group_gemm", "bmm_rcr_n1")) or "softmax" in op_name:
                continue

            # Extract alignment var names and padding lengths.
            alignment_var_to_padding_length = {}
            dim_info_dict = op._extract_dims()
            for i, tensor in enumerate(op._attrs["inputs"]):
                alignment_var = _extract_mnk_name(
                    dim_info_dict, Source.INPUT, i, len(tensor._attrs["shape"]) - 1
                )
                if alignment_var is None:
                    # No alignment var is extracted. Skip padding.
                    continue
                alignment_dim = tensor._attrs["shape"][-1]
                if not isinstance(alignment_dim, IntImm):
                    raise NotImplementedError(
                        "Gemm does not support dynamic alignment dimensions "
                        "(i.e. alignment==1)! Gemm: {}".format(op)
                    )
                padding_length = get_padding_length(
                    alignment_dim.value(), tensor.dtype()
                )
                if padding_length > 0:
                    alignment_var_to_padding_length[alignment_var] = padding_length
            if len(alignment_var_to_padding_length) == 0:
                # No padding is necessary.
                continue

            _LOGGER.debug(
                "**** Apply padding ****, alignment_var_to_padding_length: \n {} \n".format(
                    alignment_var_to_padding_length
                ),
            )
            original_op_debug_str = str(op)

            # Pad A and B.
            for tensor_idx, _ in enumerate(op._attrs["inputs"][:2]):
                _pad_input_tensor(
                    op,
                    tensor_idx,
                    lambda tensor_idx, dim_idx: _extract_mnk_name(
                        dim_info_dict, Source.INPUT, tensor_idx, dim_idx
                    ),
                    alignment_var_to_padding_length,
                    new_tensor_list,
                )

            # Pad bias and extra sources if necessary.
            for tensor_idx, tensor in enumerate(op._attrs["inputs"][2:]):
                _pad_input_tensor(
                    op,
                    tensor_idx + 2,  # skip A and B
                    lambda _, dim_idx: _extract_mnk_name(
                        dim_info_dict,
                        Source.OUTPUT,  # bias alignment follows output alignment
                        0,  # always check output[0]
                        dim_idx
                        + len(op._attrs["outputs"][0]._attrs["shape"])
                        - len(tensor._attrs["shape"]),  # handle bias broadcast case
                    ),
                    alignment_var_to_padding_length,
                    new_tensor_list,
                )

            # Replaces the old op with the new op.
            for tensor_input in op._attrs["inputs"]:
                tensor_input._attrs["dst_ops"].discard(op)
            new_op = type(op)(**op._get_op_attributes())
            new_op._attrs["split_k"] = op._attrs["split_k"]
            if "alpha" in op._attrs:
                new_op._attrs["alpha"] = op._attrs["alpha"]
            new_output = new_op(*op._attrs["inputs"])
            new_tensor_list.append(new_output)
            original_output = op._attrs["outputs"][0]
            transform_utils.copy_tensor_attributes(new_output, original_output)

            # Slice output if necessary.
            new_output = _slice_output_tensor(
                new_output, original_output, new_tensor_list
            )
            transform_utils.replace_tensor(original_output, new_output)
            transform_utils.remove_tensor_from_sorted_graph(original_output)

            _LOGGER.debug(
                "**** Apply padding ****, replace op \n {} \n with \n {} \n".format(
                    original_op_debug_str, new_op
                ),
            )

        new_sorted_graph.extend(new_tensor_list)

    new_sorted_graph = transform_utils.sanitize_sorted_graph(new_sorted_graph)

    return new_sorted_graph

    Applies padding to gemms to use SM80 kernels.
    SM80 kernels require min_alignment == 2.
    
python/aitemplate/compiler/transform/apply_padding.py