def meets_special_kernel_conditions(func_attrs, input_type, output_type) -> bool:
    """
    The reduction op must meet all conditions to use the special kernel:
    1. The input and output types are the same.
    2. The input and output memory are contiguous.
    3. The tensor is rank-3 and the reduction_dim=1.
    4. The reduction dimension length is less than 10 and the # of threads needed
    to do vectorize loads/stores is less than 1024.
    """
    if input_type != output_type:
        return False

    input_accessors = func_attrs["output_accessors"]
    output_accessors = func_attrs["output_accessors"]
    if not (
        output_accessors
        and input_accessors
        and output_accessors[0].is_contiguous
        and input_accessors[0].is_contiguous
    ):
        return False

    input_shape: List[IntVar] = func_attrs["inputs"][0].shape()
    reduction_axes = func_attrs["reduction_axes"]
    reduction_axis = reduction_axes[0]
    if not (
        len(input_shape) == 3  # This kernel only supports rank 3 tensors.
        and len(reduction_axes) == 1
        and reduction_axis == 1
        and is_static_dimension(input_shape, reduction_axis)
        and is_static_dimension(input_shape, reduction_axis + 1)
    ):
        return False

    def _get_largest_aligned_vector(input_type, nelems) -> (str, int):
        for vtype, vbytesize in vector_types[input_type]:
            vcapacity = int(vbytesize / bytesize[input_type])
            if nelems % vcapacity == 0:
                return vtype, vcapacity

    MAX_BLOCK_DIM = 1024
    reduction_axis_len = input_shape[reduction_axis].value()
    cross_axis_len = input_shape[reduction_axis + 1].value()
    _, vector_capacity = _get_largest_aligned_vector(input_type, nelems=cross_axis_len)
    if not (
        reduction_axis_len <= 10 and cross_axis_len / vector_capacity <= MAX_BLOCK_DIM
    ):
        return False

    return True
The reduction op must meet all conditions to use the special kernel:
1. The input and output types are the same.
2. The input and output memory are contiguous.
3. The tensor is rank-3 and the reduction_dim=1.
4. The reduction dimension length is less than 10 and the # of threads needed
to do vectorize loads/stores is less than 1024.
python/aitemplate/backend/cuda/reduce/reduce_common_slim_tensor.py