# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

# Adapted from https://github.com/sgl-project/sglang/blob/4cb53ecd0cffceb6dee5c011a58f65997a86f151/python/sglang/srt/layers/quantization/int8_kernel.py
import functools
import json
import os
from typing import Any, Dict, List, Optional, Tuple

from aiter.ops.triton.utils.core import AITER_TRITON_CONFIGS_PATH
import aiter.ops.triton.utils.arch_info as arch_info
from aiter import logger

import torch
import triton
import triton.language as tl

'''
configs = [
    triton.Config({"BLOCK_SIZE": BLOCK_SIZE}, num_warps=num_warps, num_stages=num_stages)
    for BLOCK_SIZE in [2**n for n in range(7, 15)] for num_warps in [1, 2, 4, 8, 16] for num_stages in [1, 2]\
]
@triton.autotune(
    configs=configs,
    key=["M", "GROUP_SIZE"],
    perf_debug=True,
    enable=int(os.getenv("TRITON_DO_AUTOTUNING", 0)) == 1,
    config_hook_params=ConfigHookParams(
        cache=JsonCache(lambda key: get_w8a8_group_quant_config_filepath(**key)), # Save best config to this file
        key_hook=MKeyHook(), # Use "M" as the key
        #extra_config_hook=None
    ),
    #prune_configs_by={
    #    "early_config_prune": prune_configs
    #}
)
'''
@triton.jit
def _per_token_group_quant_int8(
    # Pointers to inputs and output
    y_ptr,
    y_q_ptr,
    y_s_ptr,
    M,
    # Avoid to divide zero
    eps,
    int8_min,
    int8_max,
    GROUP_SIZE: tl.constexpr,
    BLOCK_SIZE: tl.constexpr
):
    """A Triton-accelerated function to perform per-token-group
    quantization on a tensor.

    This function converts the tensor values into int8 values.
    """
    g_id = tl.program_id(0)
    y_ptr += g_id * BLOCK_SIZE
    y_q_ptr += g_id * BLOCK_SIZE
    S_NUM: tl.constexpr = BLOCK_SIZE // GROUP_SIZE
    y_s_ptr += g_id * S_NUM

    cols = tl.arange(0, BLOCK_SIZE)  # N <= BLOCK_SIZE
    s_cols = tl.arange(0, S_NUM)
    mask = g_id * BLOCK_SIZE + cols < M

    y = tl.load(y_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    y = tl.reshape(y, (S_NUM, GROUP_SIZE))
    # Quant
    _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps)
    y_s = (_absmax / int8_max).reshape(S_NUM, 1)
    y_q = tl.clamp(y / y_s, int8_min, int8_max).to(y_q_ptr.dtype.element_ty)

    y_q = tl.reshape(y_q, (S_NUM * GROUP_SIZE))
    y_s = tl.reshape(y_s, (S_NUM))

    tl.store(y_q_ptr + cols, y_q, mask=mask)
    tl.store(y_s_ptr + s_cols, y_s.to(y_s_ptr.dtype.element_ty))

@functools.lru_cache
def get_w8a8_group_quant_config_filepath(M: int, GROUP_SIZE: int) -> str:
    device_name = arch_info.get_device()
    if device_name.lower().startswith("bw"):
        device_name = "BW200"
    if "k100" in device_name.lower():
        device_name = "K100_AI"
    json_file_name = f"w8a8_per_token_group_quant_device_name={device_name},group_size={GROUP_SIZE}.json"
    config_file_path = os.path.join(
        f"{AITER_TRITON_CONFIGS_PATH}", "group_quant", json_file_name
    )
    return config_file_path

@functools.lru_cache
def get_w8a8_group_quant_configs(
    M: int, groupSize: int
) -> Optional[Dict[int, Any]]:
    """
    Return optimized configurations for the w8a8 block fp8 kernel.

    The return value will be a dictionary that maps an irregular grid of
    batch sizes to configurations of the w8a8 block fp8 kernel. To evaluate the
    kernel on a given batch size bs, the closest batch size in the grid should
    be picked and the associated configuration chosen to invoke the kernel.
    """
    config_file_path = get_w8a8_group_quant_config_filepath(M, groupSize)
    if os.path.exists(config_file_path):
        with open(config_file_path) as f:
            #logger.info(
            #    "Using configuration from %s for W8A8 GROUP QUANT kernel.",
            #    config_file_path,
            #)
            # If a configuration has been found, return it
            return {int(key): val for key, val in json.load(f).items()}

    # If no optimized configuration is available, we will use the default
    # configuration
    logger.warning(
        (
            "Using default W8A8 GROUP QUANT kernel config. Performance might "
            "be sub-optimal! Config file not found at %s"
        ),
        config_file_path,
    )
    return None

def per_token_group_quant_int8(
    x: torch.Tensor,
    group_size: int,
    eps: float = 1e-10,
    dtype: torch.dtype = torch.int8,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Function to perform per-token-group quantization on an input tensor `x`.

    It converts the tensor values into signed int8 values and returns the
    quantized tensor along with the scaling factor used for quantization.

    Args:
        x: The input tenosr with ndim >= 2.
        group_size: The group size used for quantization.
        eps: The minimum to avoid dividing zero.
        dtype: The dype of output tensor. Note that only `torch.int8`
            is supported for now.

    Returns:
        Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the
            scaling factor for quantization.
    """
    assert (x.shape[-1] % group_size == 0
            ), "the last dimension of `x` cannot be divisible by `group_size`"
    assert x.is_contiguous(), "`x` is not contiguous"

    iinfo = torch.iinfo(dtype)
    int8_max = iinfo.max
    int8_min = iinfo.min

    x_q = torch.empty_like(x, device=x.device, dtype=dtype)
    x_s = torch.empty(
        x.shape[:-1] + (x.shape[-1] // group_size, ),
        device=x.device,
        dtype=torch.float32,
    )

    M = x.numel()
    configs = get_w8a8_group_quant_configs(M, group_size)
    if configs:
        config = configs[min(configs.keys(), key=lambda x: abs(x - M))]
    else:
        config = {
            "BLOCK_SIZE": 128,
            "num_warps": 1,
            "num_stages": 1,
        }

    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE']),
    )
    _per_token_group_quant_int8[grid](
        x,
        x_q,
        x_s,
        M,
        eps,
        int8_min=int8_min,
        int8_max=int8_max,  
        GROUP_SIZE=group_size,
        **config
    )

    return x_q, x_s