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
import transformer_engine_torch as tex

from triton.language.extra import libdevice

@triton.jit
def _per_token_quant_int8(
    x_ptr,
    xq_ptr,
    scale_ptr,
    stride_x,
    stride_xq,
    N,
    BLOCK: tl.constexpr,
):
    # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
    row_id = tl.program_id(0)

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

    x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask,
                other=0.0).to(tl.float32)
    absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
    scale_x = absmax / 127
    x_q = x * (127 / absmax)
    x_q = libdevice.nearbyint(x_q).to(tl.int8)

    tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask)
    tl.store(scale_ptr + row_id, scale_x)


def per_token_quant_int8(x):
    M = x.numel() // x.shape[-1]
    N = x.shape[-1]
    x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
    scales = torch.empty(x.shape[:-1] + (1, ),
                         device=x.device,
                         dtype=torch.float32)
    BLOCK = triton.next_power_of_2(N)
    # heuristics for number of warps
    num_warps = min(max(BLOCK // 256, 1), 8)

    assert x.is_contiguous()
    _per_token_quant_int8[(M, )](
        x,
        x_q,
        scales,
        stride_x=x.stride(-2),
        stride_xq=x_q.stride(-2),
        N=N,
        BLOCK=BLOCK,
        num_warps=num_warps,
        num_stages=1,
    )

    return x_q, scales

@triton.jit
def _per_token_quant_int8_v2(
    x_ptr,
    xq_ptr,
    scale_ptr,
    stride_x,
    stride_xq,
    M,
    BLOCK: tl.constexpr,
):
    col_id = tl.program_id(0)

    rows = tl.arange(0, BLOCK)
    mask = rows < M

    x = tl.load(x_ptr + rows * stride_x + col_id, mask=mask,
                other=0.0).to(tl.float32)
    absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
    scale_x = absmax / 127
    x_q = x * (127 / absmax)
    x_q = libdevice.nearbyint(x_q).to(tl.int8)

    tl.store(xq_ptr + rows * stride_xq + col_id, x_q, mask=mask)
    tl.store(scale_ptr + col_id, scale_x)


def per_token_quant_int8_v2(x):
    assert x.is_contiguous()
    x = x.view(-1, x.shape[-1])
    M = x.numel() // x.shape[-1]
    N = x.shape[-1]
    x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
    scales = torch.empty((1, x.shape[-1]),
                         device=x.device,
                         dtype=torch.float32)
    BLOCK = triton.next_power_of_2(M)
    # heuristics for number of warps
    num_warps = min(max(BLOCK // 256, 1), 8)

    _per_token_quant_int8_v2[(N, )](
        x,
        x_q,
        scales,
        stride_x=x.stride(-2),
        stride_xq=x_q.stride(-2),
        M=M,
        BLOCK=BLOCK,
        num_warps=num_warps,
        num_stages=1,
    )

    return x_q, scales



@triton.jit
def _tensorwise_dequantize_impl(
    x_ptr,
    y_ptr,
    scaleA_ptr,
    scaleB_ptr,
    stride_x,
    stride_y,
    N,
    is_add: tl.constexpr,
    is_float: tl.constexpr,
    BLOCK: tl.constexpr,
):
    row_id = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < N

    a_scale = tl.load(scaleA_ptr)
    b_scale = tl.load(scaleB_ptr)
    x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask,
                other=0.0).to(tl.float32)

    result = x * a_scale * b_scale

    if is_add:
        y = tl.load(y_ptr + row_id * stride_y + cols, mask=mask, other=0.0).to(tl.float32)
        result += y

    if is_float:
        tl.store(y_ptr + row_id * stride_y + cols, result, mask=mask)
    else:
        tl.store(y_ptr + row_id * stride_y + cols, result.to(tl.bfloat16), mask=mask)


def _tensorwise_dequantize(a_scale, b_scale, x, y, is_add=False, is_float=False):
    assert x.is_contiguous()
    x = x.view(-1, x.shape[-1])
    M = x.numel() // x.shape[-1]
    N = x.shape[-1]
    
    BLOCK = triton.next_power_of_2(N)
    # heuristics for number of warps
    num_warps = min(max(BLOCK // 256, 1), 8)

    assert x.is_contiguous()
    _tensorwise_dequantize_impl[(M, )](
        x,
        y,
        a_scale,
        b_scale,
        stride_x=x.stride(-2),
        stride_y=y.stride(-2),
        N=N,
        is_add=is_add,
        is_float=is_float,
        BLOCK=BLOCK,
        num_warps=num_warps,
        num_stages=1,
    )



@triton.jit
def _per_token_quant_fp8_to_int8(
    x_ptr,
    xq_ptr,
    scale_ptr,
    stride_x,
    stride_xq,
    N,
    fp8_scale_inv,
    BLOCK: tl.constexpr,
):
    # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
    row_id = tl.program_id(0)

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

    x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask,
                other=0.0).to(tl.float32)
    fp8_scale = tl.load(fp8_scale_inv)
    x = x * fp8_scale
    absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
    scale_x = absmax / 127
    x_q = x * (127 / absmax)
    x_q = libdevice.nearbyint(x_q).to(tl.int8)

    tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask)
    tl.store(scale_ptr + row_id, scale_x)


def per_token_quant_fp8_to_int8(x, fp8_scale_inv, inplace=False):
    assert x.is_contiguous()
    x = x.view(-1, x.shape[-1])
    M = x.numel() // x.shape[-1]
    N = x.shape[-1]
    if inplace:
        x_q = x.view(dtype=torch.int8)
    else:
        x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
    scales = torch.empty(x.shape[:-1] + (1, ),
                         device=x.device,
                         dtype=torch.float32)
    BLOCK = triton.next_power_of_2(N)
    # heuristics for number of warps
    num_warps = min(max(BLOCK // 256, 1), 8)

    assert x.is_contiguous()
    _per_token_quant_fp8_to_int8[(M, )](
        x,
        x_q,
        scales,
        stride_x=x.stride(-2),
        stride_xq=x_q.stride(-2),
        N=N,
        BLOCK=BLOCK,
        fp8_scale_inv=fp8_scale_inv,
        num_warps=num_warps,
        num_stages=1,
    )

    return x_q, scales

@triton.jit
def _per_token_quant_fp8_to_int8_opt(
    x_ptr,
    xq_ptr,
    scale_ptr,
    stride_x,
    stride_xq,
    N,
    fp8_scale_inv,
    BLOCK: tl.constexpr,
):
    # Adapted from https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
    row_id = tl.program_id(0)

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

    x = tl.load(x_ptr + row_id * stride_x + cols, mask=mask,
                other=0.0).to(tl.float32)
    scales = tl.load(scale_ptr + cols, mask=mask,
                other=0.0).to(tl.float32)
    fp8_scale = tl.load(fp8_scale_inv)
    x = x * fp8_scale
    x_q = x / scales
    x_q = libdevice.nearbyint(x_q).to(tl.int8)

    tl.store(xq_ptr + row_id * stride_xq + cols, x_q, mask=mask)


def per_token_quant_fp8_to_int8_opt(x, fp8_scale_inv, inplace=False):
    assert x.is_contiguous()
    x = x.view(-1, x.shape[-1])
    M = x.numel() // x.shape[-1]
    N = x.shape[-1]
    if inplace:
        x_q = x.view(dtype=torch.int8)
    else:
        x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
    scales = torch.empty((1, x.shape[-1]),
                         device=x.device,
                         dtype=torch.float32)
    
    tex.compute_channel_colwise_amax(x, scales, fp8_scale_inv)

    BLOCK = triton.next_power_of_2(N)
    # heuristics for number of warps
    num_warps = min(max(BLOCK // 256, 1), 8)

    _per_token_quant_fp8_to_int8_opt[(M, )](
        x,
        x_q,
        scales,
        stride_x=x.stride(-2),
        stride_xq=x_q.stride(-2),
        N=N,
        BLOCK=BLOCK,
        fp8_scale_inv=fp8_scale_inv,
        num_warps=num_warps,
        num_stages=1,
    )

    return x_q, scales



@triton.jit
def _per_token_quant_fp8_to_int8_v2(
    x_ptr,
    xq_ptr,
    scale_ptr,
    stride_x,
    stride_xq,
    M,
    fp8_scale_inv,
    BLOCK: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    col_id = tl.program_id(0)

    rows = tl.arange(0, BLOCK)
    cols = tl.arange(0, BLOCK_N)
    offset_cols = col_id * BLOCK_N + cols

    mask = rows[:,None] < M

    x = tl.load(x_ptr + rows[:,None] * stride_x + offset_cols[None,:], mask=mask, other=0.0).to(tl.float32)
    fp8_scale = tl.load(fp8_scale_inv)
    x = x * fp8_scale
    absmax = tl.maximum(tl.max(tl.abs(x)), 1e-10)
    scale_x = absmax / 127
    x_q = x * (127 / absmax)
    x_q = libdevice.nearbyint(x_q).to(tl.int8)

    tl.store(xq_ptr + rows[:,None] * stride_xq + offset_cols[None,:], x_q, mask=mask)
    tl.store(scale_ptr + offset_cols, scale_x)


def per_token_quant_fp8_to_int8_v2(x, fp8_scale_inv, inplace=False):
    assert x.is_contiguous()
    x = x.view(-1, x.shape[-1])
    M = x.numel() // x.shape[-1]
    N = x.shape[-1]
    if inplace:
        x_q = x.view(dtype=torch.int8)
    else:
        x_q = torch.empty_like(x, device=x.device, dtype=torch.int8)
    scales = torch.empty((1, x.shape[-1]),
                         device=x.device,
                         dtype=torch.float32)
    BLOCK = triton.next_power_of_2(M)
    # heuristics for number of warps
    num_warps = min(max(BLOCK // 256, 1), 8)

    _per_token_quant_fp8_to_int8_v2[(N//32, )](
        x,
        x_q,
        scales,
        stride_x=x.stride(-2),
        stride_xq=x_q.stride(-2),
        M=M,
        fp8_scale_inv=fp8_scale_inv,
        BLOCK=BLOCK,
        BLOCK_N=triton.next_power_of_2(32),
        num_warps=num_warps,
        num_stages=1,
    )

    return x_q, scales

@torch.compile(mode="max-autotune-no-cudagraphs")
def channelwise_dequantize(A, B, C):
    out_scales = A * B
    return (out_scales * C.to(dtype=torch.float32)).to(torch.bfloat16)

@torch.compile(mode="max-autotune-no-cudagraphs")
def channelwise_dequantize_transA(A, B, C):
    out_scales = A.T * B
    return (out_scales * C.to(dtype=torch.float32)).to(torch.bfloat16)

@torch.compile(mode="max-autotune-no-cudagraphs")
def channelwise_dequantize_transA_float(A, B, C):
    out_scales = A.T * B
    return out_scales * C.to(dtype=torch.float32)

@torch.compile(mode="max-autotune-no-cudagraphs")
def channelwise_dequantize_transA_add(A, B, C, D):
    out_scales = A.T * B
    D.add_((out_scales * C.to(dtype=torch.float32)).to(torch.bfloat16))

@torch.compile(mode="max-autotune-no-cudagraphs")
def channelwise_dequantize_transA_float_add(A, B, C, D):
    out_scales = A.T * B
    D.add_(out_scales * C.to(dtype=torch.float32))

@torch.compile(mode="max-autotune-no-cudagraphs")
def channelwise_dequantize_transB(A, B, C):
    out_scales = A * B.T
    return (out_scales * C.to(dtype=torch.float32)).to(torch.bfloat16)

def tensorwise_dequantize(A, B, C, D):
    _tensorwise_dequantize(A, B, C, D, is_add=False, is_float=False)

def tensorwise_dequantize_float(A, B, C, D):
    _tensorwise_dequantize(A, B, C, D, is_add=False, is_float=True)

def tensorwise_dequantize_add(A, B, C, D):
    _tensorwise_dequantize(A, B, C, D, is_add=True, is_float=False)

def tensorwise_dequantize_float_add(A, B, C, D):
    _tensorwise_dequantize(A, B, C, D, is_add=True, is_float=True)

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()




