def filter_cutlass_3x_ops(op_instance, func_attrs):
    """Filter out CUTLASS 3.x ops with incompatible alignment requirements.

    The CUTLASS 3.x ops have stricter alignment requirements compared to
    the CUTLASS 2.x ops (due to TMA). These alignment requirements are used
    to initially filter them out in the `function_filter` below. However, the
    required alignments of the GEMM op inputs and outputs may change due to
    TensorAccessor-related optimizations, which are introduced to the model
    graph *after* the initial filtering.

    In this function, the (possible) TA-related alignment updates are checked
    once again and the CUTLASS 3.x ops not satisfying these requirements are
    filtered out. Importantly, due to input/output alignment flexibilit of the
    CUTLASS 2.x ops, their alignment requirements are corrected using the
    TA-imposed alignments in the `update_alignments_in_gemm_instance` function
    above. But this correction is not possible for the CUTLASS 3.x ops, as they
    won't work with the lower alignment values. That's why the CUTLASS 3.x ops
    are filtered out by this function in such cases.
    """
    import cutlass_lib

    a_alignment, b_alignment, epilogue_alignment = get_tensor_accessor_alignments(
        func_attrs
    )

    result_2x, result_3x = {}, {}
    for op_name, op in op_instance.items():
        if op.gemm_kind == cutlass_lib.library.GemmKind.Universal3x:
            if (
                op.A.alignment <= a_alignment
                and op.B.alignment <= b_alignment
                and op.C.alignment <= epilogue_alignment
            ):
                result_3x[op_name] = op
        else:
            result_2x[op_name] = op

    has_ops_with_tma_epilogue = False
    if result_3x:
        for op in result_3x.values():
            if has_tma_epilogue(op):
                has_ops_with_tma_epilogue = True
                break

        if has_ops_with_tma_epilogue:
            # when there are ops with TMA epilogue, keep only those
            # for better performance / shorter profiler compilation time
            result_3x = {
                op_name: op for op_name, op in result_3x.items() if has_tma_epilogue(op)
            }

    return {
        # CUTLASS 3.x kernels can cause power throttling:
        # we want to generate the 2.x kernels first to avoid
        # performance side effects caused by the 3.x kernels
        **result_2x,
        **result_3x,
    }, has_ops_with_tma_epilogue
Filter out CUTLASS 3.x ops with incompatible alignment requirements.

The CUTLASS 3.x ops have stricter alignment requirements compared to
the CUTLASS 2.x ops (due to TMA). These alignment requirements are used
to initially filter them out in the `function_filter` below. However, the
required alignments of the GEMM op inputs and outputs may change due to
TensorAccessor-related optimizations, which are introduced to the model
graph *after* the initial filtering.

In this function, the (possible) TA-related alignment updates are checked
once again and the CUTLASS 3.x ops not satisfying these requirements are
filtered out. Importantly, due to input/output alignment flexibilit of the
CUTLASS 2.x ops, their alignment requirements are corrected using the
TA-imposed alignments in the `update_alignments_in_gemm_instance` function
above. But this correction is not possible for the CUTLASS 3.x ops, as they
won't work with the lower alignment values. That's why the CUTLASS 3.x ops
are filtered out by this function in such cases.
python/aitemplate/backend/cuda/gemm_universal/common.py