import torch
import time
from typing import Optional, Type,Any, Dict, List, Tuple
import pandas as pd
import os
import json
import triton
import triton.language as tl
import pandas as pd
import logging
import math
from transformer_engine.pytorch.fp8 import blockwise_fp8_block_len


def to_int8(tensor: torch.Tensor):
    return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)

@triton.jit
def _per_token_group_quant_int8(
    # Pointers to inputs and output
    y_ptr,
    y_q_ptr,
    y_s_ptr,
    # Stride of input
    y_stride,
    # Collums of input
    N,
    # Avoid to divide zero
    eps,
    # Information for int8
    int8_min,
    int8_max,
    # Meta-parameters
    BLOCK: tl.constexpr,
):
    """A Triton-accelerated function to perform
    per-token-group quantization on a tensor.
    This function converts the tensor values into int8 values.
    """
    # Map the program id to the row of X and Y it should compute.
    g_id = tl.program_id(0)
    y_ptr += g_id * y_stride
    y_q_ptr += g_id * y_stride
    y_s_ptr += g_id

    cols = tl.arange(0, BLOCK)  # N <= BLOCK
    mask = cols < N

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

    tl.store(y_q_ptr + cols, y_q, mask=mask)
    tl.store(y_s_ptr, y_s)


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)
    M = x.numel() // group_size
    N = group_size
    x_s = torch.empty(
        x.shape[:-1] + (x.shape[-1] // group_size,),
        device=x.device,
        dtype=torch.float32,
    )

    BLOCK = triton.next_power_of_2(N)#N是blocksize[1]
    # heuristics for number of warps
    num_warps = min(max(BLOCK // 256, 1), 8)
    num_stages = 1
  
    _per_token_group_quant_int8[(M,)](
        x,
        x_q,
        x_s,
        group_size,
        N,
        eps,
        int8_min=int8_min,
        int8_max=int8_max,
        BLOCK=BLOCK,
        num_warps=num_warps,
        num_stages=num_stages,
    )

    return x_q, x_s,BLOCK,num_warps,num_stages,M

def _int8_gemm_helper(m: int,
                             n: int,
                             k: int,
                             out_dtype: Type[torch.dtype] = torch.float16,
                             device: str = "cuda",
                             block_size: List[int]=[blockwise_fp8_block_len,blockwise_fp8_block_len],
                             best_config:Optional[list] = None):
    # Test for a cutlass kernel with per-token activation quantization
    # and per-output channel weight quantization.
    input = (torch.randn((m, k), device=device) * 5).to(dtype=out_dtype)
    weight = to_int8(torch.randn((n ,k), device=device) * 5)

    weight_scale = (torch.randn((math.ceil(n/block_size[0]), math.ceil(k/block_size[1])), device=device,
                           dtype=torch.float32))
    print("input.dtype:",input.dtype)
    #print("m:{} n:{} k:{},weight_scale.shape:{}".format(m,n,k,weight_scale.shape))

    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[0]]
    
    q_input, x_scale,_,_,_,_ = per_token_group_quant_int8(input_2d, block_size[1])
    
    return q_input, x_scale,weight,weight_scale


def _int8_gemm_helper_b(m: int,
                             n: int,
                             k: int,
                             out_dtype: Type[torch.dtype] = torch.float16,
                             device: str = "cuda",
                             block_size: List[int]=[blockwise_fp8_block_len,blockwise_fp8_block_len],
                             best_config:Optional[list] = None):
    # Test for a cutlass kernel with per-token activation quantization
    # and per-output channel weight quantization.
    input = (torch.randn((m, k), device=device) * 5).to(dtype=out_dtype)
    weight = to_int8(torch.randn((n ,k), device=device) * 5)

    weight_scale = (torch.randn((n, math.ceil(k/block_size[1])), device=device,
                           dtype=torch.float32))
    print("input.dtype:",input.dtype)
    #print("m:{} n:{} k:{},weight_scale.shape:{}".format(m,n,k,weight_scale.shape))

    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[0]]
    
    q_input, x_scale,_,_,_,_ = per_token_group_quant_int8(input_2d, block_size[1])
    
    return q_input, x_scale,weight,weight_scale


def _int8_gemm_helper_test(m: int,
                             n: int,
                             k: int,
                             out_dtype: Type[torch.dtype] = torch.float16,
                             device: str = "cuda",
                             block_size: List[int]=[blockwise_fp8_block_len,blockwise_fp8_block_len],
                             best_config:Optional[list] = None):
    # Test for a cutlass kernel with per-token activation quantization
    # and per-output channel weight quantization.
    input = (torch.randn((m, k), device=device) * 5).to(dtype=out_dtype)
    weight = (torch.randn((n ,k), device=device) * 5).t().to(dtype=out_dtype)

    print("input.dtype:",input.dtype)
    input_2d = input.view(-1, input.shape[-1])
    output_shape = [*input.shape[:-1], weight.shape[0]]
    
    q_input, x_scale,BLOCK,num_warps,num_stages,M= per_token_group_quant_int8(input_2d, block_size[1])
    
    start_time_ = time.time()  # 开始计时
    for it in range(1000):
        q_input, x_scale,_,_,_,_ =per_token_group_quant_int8(input_2d, block_size[1])
    torch.cuda.synchronize()    
    end_time_ = time.time()  # 结束计时
    
    elapsed_time = round((end_time_ - start_time_) *1000 ,7)# 计算耗时  
    print("_time:{} us\n".format(elapsed_time))
    
    return q_input, x_scale,elapsed_time,BLOCK,num_warps,num_stages,M


def main():
    m_list=[1,2,4,8,16,32,64,128,256,512,1024,2048,4096,8192,16384,32768]
    n_list=[576,2048,7168,256,7168,1536,1536]
    k_list=[7168,512,1024,7168,128,7168,1536]
    block_size=[blockwise_fp8_block_len,blockwise_fp8_block_len]
    out_dtype=torch.bfloat16
    _n=[]
    _k=[]
    _m=[]
    
    
    config_blocks=[]
    config_num_warps=[]
    config_num_stages=[]
    config_M=[]
    
    cost_times=[]
    for i in range(0,len(k_list),1):
        for m in m_list:
            print("m:{} n:{} k:{} ".format(m,n_list[i],k_list[i]))
    
            q_input, x_scale,elapsed_time,BLOCK,num_warps,num_stages,M=_int8_gemm_helper_test(m=m,n=n_list[i],k=k_list[i],block_size=block_size,out_dtype=torch.bfloat16)

            cost_times.append(elapsed_time)
            _n.append(n_list[i])
            _k.append(k_list[i])
            _m.append(m)
            config_blocks.append(BLOCK)
            config_num_warps.append(num_warps)
            config_num_stages.append(num_stages)
            config_M.append(M)
    
    # 创建一个包含这三个列表的 DataFrame
    df = pd.DataFrame({'m':_m,'n':_n,'k':_k,'量化算子耗时': cost_times,'BLOCK':config_blocks,'num_warps':config_num_warps,'config_num_stages':config_num_stages,'config_M':config_M})
    # 将 DataFrame 写入 Excel 文件
    df.to_excel('output.xlsx', index=False)

    print("表格已保存到 output.xlsx 文件中。")    

if __name__ == "__main__":
    main()




