#
#  Copyright (2024) Hewlett Packard Enterprise Development LP
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  You may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#  http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

"""Helper functionality."""

from math import ceil

from ..types import LaunchConfiguration


def simple_kernel_dimensions(
    threads_needed: int, max_threads_per_block: int
) -> LaunchConfiguration:
    """
    Determines a simple one-dimensional set of CUDA kernel dimensions,
    i.e. `(dim_grid, dim_block)` based on a wanted amount of threads
    """

    if threads_needed <= 0:
        return ((0, 0, 0), (0, 0, 0))

    threads_per_block = min(threads_needed, max_threads_per_block)
    blocks_per_grid = ceil(threads_needed / threads_per_block)

    dim_grid = (blocks_per_grid, 1, 1)
    dim_block = (threads_per_block, 1, 1)

    return dim_grid, dim_block
