    def _compute_ws_size_strided(
        self,
        extent: List[int],
        reduction_axis: int,
        vector_length: int,
        accumulation_type: str,
    ) -> int:
        """
        Compute workspace size for contiguous reduction kernels.
        Reference: cutlass reduction/device/tensor_reduce_affine_strided.h
        """
        k_rank = 4
        k_reduced_rank = 3
        k_inner_rank = k_rank - k_reduced_rank
        k_threads = 256
        target_threadblock_count = 128

        outer_count = 1
        for dim_val in extent[: k_reduced_rank - 1]:
            outer_count *= dim_val
        inner_count = 1
        for dim_val in extent[k_reduced_rank - 1 : k_reduced_rank - 1 + k_inner_rank]:
            inner_count *= dim_val
        extent_c = extent[k_rank - 1]
        vectors_c = (extent_c - 1 + vector_length) // vector_length
        cta_width = k_threads * vector_length

        def _reshape_pow2(ext, count):
            if ext > count:
                return 1
            x = 1
            while count >= ext * 2:
                count >>= 1
                x <<= 1
            return x

        cta_ways = _reshape_pow2(extent_c, cta_width)
        cta_threads_x = k_threads / cta_ways
        threadblock_shape_y = 1
        threadblock_shape_z = min(cta_ways, 64)
        cta_count_x = (vectors_c + cta_threads_x - 1) // cta_threads_x
        cta_count_y = max(1, target_threadblock_count // cta_count_x)
        if cta_count_y * threadblock_shape_y > outer_count:
            cta_count_y = (outer_count + threadblock_shape_y - 1) // threadblock_shape_y
        cta_count_z = max(1, target_threadblock_count / cta_count_y)
        if cta_count_z * threadblock_shape_z > inner_count:
            cta_count_z = (inner_count + threadblock_shape_z - 1) // threadblock_shape_z
        if int(cta_count_z) == 1:
            return 0
        vector_size_bytes = vector_length * get_dtype_size(accumulation_type)
        workspace_stride = extent[k_rank - 1] * vector_size_bytes
        return workspace_stride * outer_count * cta_count_z
Compute workspace size for contiguous reduction kernels.
Reference: cutlass reduction/device/tensor_reduce_affine_strided.h
python/aitemplate/compiler/ops/reduce/reduce_common.py