# Copyright (c) 2024, NVIDIA CORPORATION & AFFILIATES. ALL RIGHTS RESERVED.
#
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception


import numba

from .._common import (
    make_binary_tempfile,
    normalize_dim_param,
    normalize_dtype_param,
)
from .._types import (
    Algorithm,
    Dependency,
    DependentArray,
    DependentPointer,
    Invocable,
    Pointer,
    TemplateParameter,
)

CUB_BLOCK_LOAD_ALGOS = {
    "direct": "::cub::BLOCK_LOAD_DIRECT",
    "striped": "::cub::BLOCK_LOAD_STRIPED",
    "vectorize": "::cub::BLOCK_LOAD_VECTORIZE",
    "transpose": "::cub::BLOCK_LOAD_TRANSPOSE",
    "warp_transpose": "::cub::BLOCK_LOAD_WARP_TRANSPOSE",
    "warp_transpose_timesliced": "::cub::BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED",
}

CUB_BLOCK_STORE_ALGOS = {
    "direct": "::cub::BLOCK_STORE_DIRECT",
    "striped": "::cub::BLOCK_STORE_STRIPED",
    "vectorize": "::cub::BLOCK_STORE_VECTORIZE",
    "transpose": "::cub::BLOCK_STORE_TRANSPOSE",
    "warp_transpose": "::cub::BLOCK_STORE_WARP_TRANSPOSE",
    "warp_transpose_timesliced": "::cub::BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED",
}


def make_load(dtype, threads_per_block, items_per_thread=1, algorithm="direct"):
    """
    Creates a block-wide load primitive.

    Returns a callable object that can be linked to and invoked from
    device code. It can be invoked with the following signatures:

    - `(src: numba.types.Array, dest: numba.types.Array) -> None`:
      Each thread loads `items_per_thread` items from `src` into `dest`.
      `dest` must contain at least `items_per_thread` items.

    Different data movement strategies can be selected via the
    `algorithm` parameter:

    - `algorithm="direct"` (default): Reads blocked data directly.
    - `algorithm="striped"`: Reads striped data directly.
    - `algorithm="vectorize"`: Reads blocked data directly using
      CUDA's built-in vectorized loads as a coalescing optimization.
    - `algorithm="transpose"`: Reads striped data and then locally
      transposes it into a blocked arrangement.
    - `algorithm="warp_transpose"`: Reads warp-striped data and then
      locally transposes it into a blocked arrangement.
    - `algorithm="warp_transpose_timesliced"`: Reads warp-striped data
      and then locally transposes it into a blocked arrangement one
      warp at a time.

    For more details, read the corresponding CUB C++ documentation:
    https://nvidia.github.io/cccl/cub/api/classcub_1_1BlockLoad.html

    Args:
        dtype: Data type being loaded
        threads_per_block: Number of threads in a block. Can be an
            integer or a tuple of 2 or 3 integers.
        items_per_thread: The number of items each thread loads
        algorithm: The data movement algorithm to use

    Example:
        The code snippet below illustrates a striped load and store of
        128 integer items by 32 threads, with each thread handling
        4 integers.

        .. literalinclude:: ../../python/cuda_cccl/tests/coop/test_block_load_store_api.py
            :language: python
            :dedent:
            :start-after: example-begin imports
            :end-before: example-end imports

        The following snippet shows how to invoke the returned
        ``block_load`` and ``block_store`` primitives:

        .. literalinclude:: ../../python/cuda_cccl/tests/coop/test_block_load_store_api.py
            :language: python
            :dedent:
            :start-after: example-begin load_store
            :end-before: example-end load_store
    """
    dim = normalize_dim_param(threads_per_block)
    dtype = normalize_dtype_param(dtype)

    template = Algorithm(
        "BlockLoad",
        "Load",
        "block_load",
        ["cub/block/block_load.cuh"],
        [
            TemplateParameter("T"),
            TemplateParameter("BLOCK_DIM_X"),
            TemplateParameter("ITEMS_PER_THREAD"),
            TemplateParameter("ALGORITHM"),
            TemplateParameter("BLOCK_DIM_Y"),
            TemplateParameter("BLOCK_DIM_Z"),
        ],
        [
            [
                Pointer(numba.uint8),
                DependentPointer(Dependency("T")),
                DependentArray(Dependency("T"), Dependency("ITEMS_PER_THREAD")),
            ]
        ],
    )
    specialization = template.specialize(
        {
            "T": dtype,
            "BLOCK_DIM_X": dim[0],
            "ITEMS_PER_THREAD": items_per_thread,
            "ALGORITHM": CUB_BLOCK_LOAD_ALGOS[algorithm],
            "BLOCK_DIM_Y": dim[1],
            "BLOCK_DIM_Z": dim[2],
        }
    )
    return Invocable(
        temp_files=[
            make_binary_tempfile(ltoir, ".ltoir")
            for ltoir in specialization.get_lto_ir()
        ],
        temp_storage_bytes=specialization.temp_storage_bytes,
        temp_storage_alignment=specialization.temp_storage_alignment,
        algorithm=specialization,
    )


def make_store(dtype, threads_per_block, items_per_thread=1, algorithm="direct"):
    """
    Creates a block-wide store primitive.

    Returns a callable object that can be linked to and invoked from
    device code. It can be invoked with the following signatures:

    - `(dest: numba.types.Array, src: numba.types.Array) -> None`:
      Each thread stores `items_per_thread` items from `src` into `dest`.
      `src` must contain at least `items_per_thread` items.

    Different data movement strategies can be selected via the
    `algorithm` parameter:

    - `algorithm="direct"` (default): Writes blocked data directly.
    - `algorithm="striped"`: Writes striped data directly.
    - `algorithm="vectorize"`: Writes blocked data directly using
      CUDA's built-in vectorized stores as a coalescing optimization.
    - `algorithm="transpose"`: Locally transposes blocked data into a
      striped arrangement before writing to memory.
    - `algorithm="warp_transpose"`: Locally transposes blocked data
      into a warp-striped arrangement before writing to memory.
    - `algorithm="warp_transpose_timesliced"`: Locally transposes
      blocked data into a warp-striped arrangement before writing to
      memory. To reduce shared memory requirements, only one warp's
      worth of shared memory is provisioned and time-sliced across warps.

    For more details, read the corresponding CUB C++ documentation:
    https://nvidia.github.io/cccl/cub/api/classcub_1_1BlockStore.html

    Args:
        dtype: Data type being stored
        threads_per_block: Number of threads in a block. Can be an
            integer or a tuple of 2 or 3 integers.
        items_per_thread: The number of items each thread stores
        algorithm: The data movement algorithm to use

    Example:
        The code snippet below illustrates a striped load and store of
        128 integer items by 32 threads, with each thread handling
        4 integers.

        .. literalinclude:: ../../python/cuda_cccl/tests/coop/test_block_load_store_api.py
            :language: python
            :dedent:
            :start-after: example-begin imports
            :end-before: example-end imports

        The following snippet shows how to invoke the returned
        ``block_load`` and ``block_store`` primitives:

        .. literalinclude:: ../../python/cuda_cccl/tests/coop/test_block_load_store_api.py
            :language: python
            :dedent:
            :start-after: example-begin load_store
            :end-before: example-end load_store
    """
    dim = normalize_dim_param(threads_per_block)
    dtype = normalize_dtype_param(dtype)

    template = Algorithm(
        "BlockStore",
        "Store",
        "block_store",
        ["cub/block/block_store.cuh"],
        [
            TemplateParameter("T"),
            TemplateParameter("BLOCK_DIM_X"),
            TemplateParameter("ITEMS_PER_THREAD"),
            TemplateParameter("ALGORITHM"),
            TemplateParameter("BLOCK_DIM_Y"),
            TemplateParameter("BLOCK_DIM_Z"),
        ],
        [
            [
                Pointer(numba.uint8),
                DependentPointer(Dependency("T")),
                DependentArray(Dependency("T"), Dependency("ITEMS_PER_THREAD")),
            ]
        ],
    )
    specialization = template.specialize(
        {
            "T": dtype,
            "BLOCK_DIM_X": dim[0],
            "ITEMS_PER_THREAD": items_per_thread,
            "ALGORITHM": CUB_BLOCK_STORE_ALGOS[algorithm],
            "BLOCK_DIM_Y": dim[1],
            "BLOCK_DIM_Z": dim[2],
        }
    )
    return Invocable(
        temp_files=[
            make_binary_tempfile(ltoir, ".ltoir")
            for ltoir in specialization.get_lto_ir()
        ],
        temp_storage_bytes=specialization.temp_storage_bytes,
        temp_storage_alignment=specialization.temp_storage_alignment,
        algorithm=specialization,
    )
