"""AWQ Triton Implementation.

This module contains the AWQ (Activation-Weight Quantization) implementation using Triton.
Cloned from vllm main branch (commit:cb080f32) and modified to fit roc.
Original file path: vllm/model_executor/layers/quantization/awq_triton.py
"""

# SPDX-License-Identifier: Apache-2.0

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

from aiter.ops.triton.gemm_a16w4 import (
    awq_dequantize_triton,
    gemm_a16w4,
    reverse_awq_order,
    AWQ_TRITON_SUPPORTED_GROUP_SIZES,
)

device = "cuda"

AWQ_GEMM_KNG_CASES = [
    # "K, N, G"
    (256, 576, 64),
    (256, 1536, 64),
    (256, 3072, 64),
    (256, 4096, 64),
    (256, 4608, 64),
    (256, 7168, 64),

    (512, 576, 64),
    (512, 1536, 64),
    (512, 3072, 64),
    (512, 4096, 64),
    (512, 4608, 64),
    (512, 7168, 64),

    (1536, 576, 64),
    (1536, 1536, 64),
    (1536, 3072, 64),
    (1536, 4096, 64),
    (1536, 4608, 64),
    (1536, 7168, 64),

    (2048, 512, 64),
    (2048, 1536, 64),
    (2048, 3072, 64),
    (2048, 4096, 64),
    (2048, 4608, 64),
    (2048, 7168, 64),

    (2304, 512, 64),
    (2304, 1536, 64),
    (2304, 3072, 64),
    (2304, 4096, 64),
    (2304, 4608, 64),
    (2304, 7168, 64),

    (7168, 512, 64),
    (7168, 1536, 64),
    (7168, 3072, 64),
    (7168, 4096, 64),
    (7168, 4608, 64),
    (7168, 7168, 64),
]

AWQ_GEMM_TEST_CASES = [
    # "M, K, N, G"
    (M, *KNG) for M in range(1, 129) for KNG in AWQ_GEMM_KNG_CASES
]
#AWQ_GEMM_TEST_CASES.append((2, 7168, 576, 64))
#AWQ_GEMM_TEST_CASES.append((10, 7168, 576, 64))
AWQ_GEMM_TEST_CASES.sort(key=lambda x: (x[0], x[1], x[2], x[3]))

# Performance benchmarking configurations
AWQ_PERF_MODEL_CASES = [
    (4096, 4096, 128),
    (4096, 8192, 128),
    (8192, 4096, 128),
]

configs = [
    triton.testing.Benchmark(
        x_names=['NUM_ROWS', 'NUM_COLS', 'GROUP_SIZE'],
        x_vals=AWQ_PERF_MODEL_CASES,
        line_arg='provider',
        line_vals=['triton', 'torch'],
        line_names=['Triton', 'Torch'],
        styles=[('red', '-'), ('blue', '--')],
        ylabel='TOPS',
        xlabel='Matrix Dimensions (Rows x Cols)',
        plot_name='AWQ Dequantize Performance',
        args={'dtype': torch.float16, 'device': device},
    )
]

@triton.testing.perf_report(configs)
def bench_awq_dequantize(NUM_ROWS, NUM_COLS, GROUP_SIZE, provider,
                        dtype=torch.float16, device="cuda"):
    """Benchmark AWQ dequantization performance."""
    warmup = 25
    rep = 10

    qweight = torch.randint(0, torch.iinfo(torch.int32).max,
                          (NUM_ROWS, NUM_COLS//8),
                          dtype=torch.int32, device=device)
    scales = torch.rand((NUM_ROWS//GROUP_SIZE, NUM_COLS),
                       dtype=dtype, device=device)
    zeros = torch.randint(0, torch.iinfo(torch.int32).max,
                         (NUM_ROWS//GROUP_SIZE, NUM_COLS//8),
                         dtype=torch.int32, device=device)

    if provider == "triton":
        fn = lambda: awq_dequantize_triton(qweight, scales, zeros,
                                         block_size_x=32, block_size_y=32)
    else:  # provider == "torch"
        fn = lambda: awq_dequantize_torch(qweight, scales, zeros, GROUP_SIZE)

    ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
    ops_per_element = 3
    total_elements = NUM_ROWS * NUM_COLS
    total_ops = total_elements * ops_per_element
    return total_ops / ms * 1e-9  # Return TOPS

def awq_dequantize_torch(
    qweight: torch.Tensor,
    scales: torch.Tensor,
    qzeros: torch.Tensor,
    group_size: int,
) -> torch.Tensor:
    """Dequantize weights using PyTorch implementation.
    
    Args:
        qweight: Quantized weight tensor
        scales: Scale factors tensor
        qzeros: Zero points tensor
        group_size: Size of groups for quantization
        
    Returns:
        Dequantized tensor
    """
    if group_size == -1:
        group_size = qweight.shape[0]

    bits = 4
    shifts = torch.arange(0, 32, bits, device=qzeros.device)

    iweights = torch.bitwise_right_shift(
        qweight[:, :, None],
        shifts[None, None, :],
    ).to(torch.int8)
    iweights = iweights.view(iweights.shape[0], -1)

    zeros = torch.bitwise_right_shift(
        qzeros[:, :, None],
        shifts[None, None, :],
    ).to(torch.int8)
    zeros = zeros.view(qzeros.shape[0], -1)
    zeros = reverse_awq_order(zeros)

    iweights = reverse_awq_order(iweights)

    iweights = torch.bitwise_and(iweights, (2**bits) - 1)
    zeros = torch.bitwise_and(zeros, (2**bits) - 1)

    scales = scales.repeat_interleave(group_size, dim=0)
    zeros = zeros.repeat_interleave(group_size, dim=0)
    return (iweights - zeros) * scales

# Benchmark configurations
bench_configs = [
    triton.testing.Benchmark(
        x_names=['M', 'K', 'N', 'G'],
        x_vals=AWQ_GEMM_TEST_CASES,
        line_arg='Provider',
        line_vals=["Triton"],
        line_names=['Triton'],
        styles=[('red', '-'), ('blue', '--')],
        ylabel='ms',
        xlabel='Matrix Dimensions (M×K×N)',
        plot_name='AWQ GEMM Performance',
        args={'device': device}
    )
]

@triton.testing.perf_report(bench_configs)
def bench_awq_gemm(M, K, N, G, Provider, device="cuda"):
    """Benchmark AWQ GEMM performance.
    
    Args:
        M: Number of rows in input matrix
        K: Number of columns in input matrix
        N: Number of columns in weight matrix (pre-quantization)
        G: AWQ group size
        provider: Implementation provider ('triton' or 'torch')
        device: Device to run on
    """
    warmup = 25
    rep = 10
    G = G if G != -1 else K

    # Generate test data
    input_tensor = torch.rand(
        (M, K),
        dtype=torch.float16,
        device=device
    )
    qweight = torch.randint(
        0,
        torch.iinfo(torch.int8).max,
        (N, K // 2),
        dtype=torch.int8,
        device=device
    )
    qzeros = torch.randint(
        0,
        torch.iinfo(torch.int8).max,
        (K // G, N // 2),
        dtype=torch.int8,
        device=device
    )
    scales = torch.rand(
        (K // G, N),
        dtype=torch.float16,
        device=device
    )

    #if provider == "triton":
    if True:
        fn = lambda: gemm_a16w4(
            input_tensor,
            qweight,
            scales,
            qzeros,
        )
        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
        
    '''
    else:  # provider == "torch"
        fn = lambda: torch.matmul(
            input_tensor, 
            awq_dequantize_torch(qweight, scales, qzeros, G)
        )
        ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
        return ms
    '''

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