# SPDX-License-Identifier: Apache-2.0

# Adapted from https://github.com/sgl-project/sglang/pull/3730
import itertools
import unittest

import torch
import triton
import triton.language as tl
import pytest
import os
import numpy as np

from aiter.ops.triton.group_quant_int8 import per_token_group_quant_int8

device = "cuda"

# For test
def native_per_token_group_quant_int8(x, group_size, eps=1e-10, dtype=torch.int8):
    """per-token-group quantization on an input tensor `x` using native torch.

    It converts the tensor values into float8 values and returns the
    quantized tensor along with the scaling factor used for quantization.
    Note that only `torch.float8_e4m3fn` is supported for now.
    """
    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_min = iinfo.min
    int8_max = iinfo.max

    x_ = x.reshape(x.numel() // group_size, group_size)
    amax = x_.abs().max(dim=-1, keepdim=True)[0].clamp(min=eps).to(torch.float32)
    x_s = amax / int8_max
    x_q = (x_ / x_s).clamp(min=int8_min, max=int8_max).to(dtype)
    x_q = x_q.reshape(x.shape)
    x_s = x_s.reshape(x.shape[:-1] + (x.shape[-1] // group_size,))

    return x_q, x_s

# Test cases for per_token_group_quant_int8
QUANT_TEST_CASES = [
    # rows, cols, group_size  # frequency
    (128, 7168, 128),  # 107326
    (128, 128, 128),  # 25812
    (1024, 128, 128),  # 25697
    (128, 1152, 128),  # 1399
    (4096, 7168, 128),  # 241
    (16384, 7168, 128),  # 241
    (512, 512, 128),  # 61
    (1024, 1024, 128),  # 61
    (1536, 1536, 128),  # 61
    (4096, 512, 128),  # 61
    (4096, 1024, 128),  # 61
    (4096, 1536, 128),  # 61
    (16384, 512, 128),  # 61
    (16384, 1024, 128),  # 61
    (16384, 1536, 128),  # 61
    (4096, 128, 128),  # 58
    (16384, 128, 128),  # 58
    (32768, 128, 128),  # 58
    (131072, 128, 128),  # 58
    (4096, 1152, 128),  # 3
    (16384, 1152, 128),  # 3
]

# Performance benchmarking configurations
quant_bench_configs = [
    triton.testing.Benchmark(
        x_names=['NUM_ROWS', 'NUM_COLS', 'GROUP_SIZE'],
        x_vals=QUANT_TEST_CASES,
        line_arg='provider',
        line_vals=['triton', 'torch'],
        line_names=['Triton', 'Torch'],
        styles=[('red', '-'), ('blue', '--')],
        ylabel='Time (ms)',
        xlabel='Matrix Dimensions (Rows x Cols)',
        plot_name='Per-Token-Group INT8 Quantization Performance',
        args={'dtype': torch.float16, 'device': 'cuda'},
    )
]

@triton.testing.perf_report(quant_bench_configs)
def bench_per_token_group_quant_int8(NUM_ROWS, NUM_COLS, GROUP_SIZE, provider,
                                   dtype=torch.float16, device="cuda"):
    """Benchmark per_token_group_quant_int8 performance."""
    warmup = 25
    rep = 10
    
    x = torch.randn((NUM_ROWS, NUM_COLS), dtype=dtype, device=device)
    
    if provider == "triton":
        fn = lambda: per_token_group_quant_int8(x, GROUP_SIZE)
    else:  # provider == "torch"
        fn = lambda: native_per_token_group_quant_int8(x, GROUP_SIZE)
    
    if int(os.getenv("TRITON_COMPILE_ONLY", 0)) == 1:
        ms = triton.testing.do_bench_compile_only(fn, warmup=warmup, rep=rep)
    else:
        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
    return ms

if __name__ == "__main__":
    os.environ["AMDGCN_USE_BUFFER_OPS"] = "1"
    bench_per_token_group_quant_int8.run(print_data=True)