def _get_read_vector_type(input_shape, input_type, force_min_vec_type=False) -> str:
    """return vector_type for reading input along reduction axis -1 (for row-major).
    In a long run, we should consider to add profiling support to reduce kernels
    and then tune vector_type. Currently, we use the following heuristics based on
    manual experiments:
    (1) for 1-d input, it's straightforward - we pickup the max vector type
        whose size is modulo of dim_val
    (2) for n-d input, if the n-2 dimension value is less than 128, we utilize
        thread-level parallelism, i.e. we don't try to read inputs across multiple
        dimensions
    (3) we ensure each thread won't read too many elements (maximum is 256 at
        the moment)

    Parameters
    ----------
    input_shape: List[IntImm]

    Returns
    -------
    str
        returns the vector type (uint4/uint2/unsigned/cutlass::half_t)
        for reading input
    """
    type_to_size_in_bit = {
        "half": 16,
        "cutlass::half_t": 16,
        "bfloat16": 16,
        "cutlass::bfloat16_t": 16,
        "float": 32,
    }

    # FIXME
    # (1) note that we don't support int8, so we don't have vector_type for one byte
    # (2) the input type is inherited from reduce_3d, so we still
    #     use cutlass::half_t for fp16. We will replace it to half once we
    #     unify our half representation
    vector_types = {
        "cutlass::half_t": [
            ("uint4", 16),
            ("uint2", 8),
            ("unsigned", 4),
            ("cutlass::half_t", 2),
        ],
        "cutlass::bfloat16_t": [
            ("uint4", 16),
            ("uint2", 8),
            ("unsigned", 4),
            ("cutlass::bfloat16_t", 2),
        ],
        "float": [
            ("uint4", 16),
            ("uint2", 8),
            ("unsigned", 4),
        ],
    }

    def _size_to_vector_type(sz_in_byte) -> str:
        """return vector_type for the given size"""
        for vec_type, sz in vector_types[input_type]:
            if sz_in_byte % sz == 0:
                return vec_type
        raise NotImplementedError("Unsupported vector size: {}".format(sz_in_byte))

    reduction_axis = -1

    if not isinstance(input_shape[reduction_axis], IntImm):
        # the last dimension is IntVar, so the best we can do in
        # terms of the read vector type is the input_type iteself
        return input_type

    rank = len(input_shape)
    reduction_dim_val = input_shape[reduction_axis]._attrs["values"][0]
    input_type_sz_in_bit = type_to_size_in_bit.get(input_type)
    assert input_type_sz_in_bit is not None and input_type_sz_in_bit % 8 == 0
    input_type_sz_in_byte = input_type_sz_in_bit / 8

    # return the minimal vector type based on the input_type
    if force_min_vec_type:
        return _size_to_vector_type(input_type_sz_in_byte)

    # 1-d tensor
    if rank == 1:
        return _size_to_vector_type(reduction_dim_val * input_type_sz_in_byte)

    # no matter input_shape[-2] is IntImm or IntVar, we get the minimal dim value
    dim_2_val = input_shape[-2]._attrs["values"][0]
    # If input_shape[-2] is too small, we utilize thread-level parallelism
    tlp_lower_bound = 128
    if dim_2_val <= tlp_lower_bound:
        return _size_to_vector_type(reduction_dim_val * input_type_sz_in_byte)
    # When dim_2_val % 2 == 1, we cannot read inputs across dimension
    if dim_2_val % 2 == 1:
        return _size_to_vector_type(reduction_dim_val * input_type_sz_in_byte)

    # Let's make sure that each thread won't read too many.
    # Currently set it to be less than 256 elements
    max_num_elems_per_thread = 256

    def _valid_vector_type(vec_type, sz_in_byte):
        if sz_in_byte % input_type_sz_in_byte != 0:
            return False
        num_elems_in_a_vec = int(sz_in_byte / input_type_sz_in_byte)
        gcd = math.gcd(reduction_dim_val, num_elems_in_a_vec)
        num_reduction_dim = int(num_elems_in_a_vec / gcd)
        if num_reduction_dim >= dim_2_val or dim_2_val % num_reduction_dim != 0:
            return False
        if num_reduction_dim == 1:
            return True
        if num_reduction_dim * reduction_dim_val > max_num_elems_per_thread:
            return False
        return True

    for vec_type, sz in vector_types[input_type]:
        if _valid_vector_type(vec_type, sz):
            return vec_type

    raise RuntimeError("Cannot find valid vector type!")
return vector_type for reading input along reduction axis -1 (for row-major).
In a long run, we should consider to add profiling support to reduce kernels
and then tune vector_type. Currently, we use the following heuristics based on
manual experiments:
(1) for 1-d input, it's straightforward - we pickup the max vector type
whose size is modulo of dim_val
(2) for n-d input, if the n-2 dimension value is less than 128, we utilize
thread-level parallelism, i.e. we don't try to read inputs across multiple
dimensions
(3) we ensure each thread won't read too many elements (maximum is 256 at
the moment)

Parameters
----------
input_shape: List[IntImm]

Returns
-------
str
returns the vector type (uint4/uint2/unsigned/cutlass::half_t)
for reading input
python/aitemplate/backend/cuda/reduce/reduce_small_axis.py