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

import torch

from vllm import _custom_ops as ops
from vllm import envs
from vllm.model_executor.layers.quantization.utils.quant_utils import (
    pack_quantized_values_into_int32,
    unpack_quantized_values_into_int32,
)
from vllm.platforms import current_platform
from vllm.scalar_type import scalar_types

from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig

_CPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8)


class CPUWNA16LinearKernel(MPLinearKernel):
    @classmethod
    def get_min_capability(cls) -> int:
        return -1

    @classmethod
    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
        if not current_platform.is_cpu():
            return False, "CPUWNA16 only supported on CPU"

        if c.weight_type not in _CPUWNA16_SUPPORTED_QUANT_TYPES:
            return (
                False,
                f"Quant type ({c.weight_type}) not supported by "
                "CPUWNA16, supported types are: "
                f"{_CPUWNA16_SUPPORTED_QUANT_TYPES}",
            )

        if c.group_size != -1 and c.group_size % 2 != 0:
            return (
                False,
                f"Group size ({c.group_size}) not supported by "
                "CPUWNA16, supported group sizes are multiples of 2",
            )

        if c.partition_weight_shape[0] % 32 != 0:
            return (
                False,
                f"Input size ({c.partition_weight_shape[0]}) not supported by "
                "CPUWNA16, supported sizes are multiples of 32",
            )

        if c.partition_weight_shape[1] % 32 != 0:
            return (
                False,
                f"Output size ({c.partition_weight_shape[1]}) not supported by "
                "CPUWNA16, supported sizes are multiples of 32",
            )

        return True, None

    # note assumes that
    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
    #  `weight_scale`  is: {input_dim = 0, output_dim = 1}
    #  `weight_zp`     is: {input_dim = 0, output_dim = 1, packed_dim = 1}
    def _process_gptq_weights_w4a16(self, layer: torch.nn.Module):
        packed_weight = getattr(layer, self.w_q_name)
        bits = self.config.weight_type.mantissa
        pack_factor = 32 // bits
        p_w_k, _ = packed_weight.size()
        input_size = p_w_k * pack_factor
        isa_hint = _get_isa_hint(getattr(layer, self.w_s_name).dtype)
        layer.isa_hint = isa_hint

        # convert input dim packed to output dim packed
        weight = unpack_quantized_values_into_int32(
            packed_weight, self.config.weight_type, 0
        )
        weight = pack_quantized_values_into_int32(weight, self.config.weight_type, 1)
        # make 16 output channel as a block and transpose to the make
        # the block contiguous
        weight = (
            weight.view(input_size, -1, 16 // pack_factor)
            .permute(1, 0, 2)
            .reshape(-1, input_size * 16 // pack_factor)
            .contiguous()
        )
        getattr(layer, self.w_q_name).data = weight

    # note assumes that
    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
    #  `weight_scale`  is: {input_dim = 0, output_dim = 1}
    #  `weight_zp`     is: {input_dim = 0, output_dim = 1, packed_dim = 1}
    def _process_gptq_weights_w4a8(self, layer: torch.nn.Module):
        packed_weight = getattr(layer, self.w_q_name)
        scales = getattr(layer, self.w_s_name)
        group_num = scales.data.size(0)
        zp_output_size = scales.data.size(1) // 8
        if self.config.zero_points:
            assert self.w_zp_name
            packed_zp = getattr(layer, self.w_zp_name)
        else:
            # w4a8 kernel always requires zp, allocate a fake zp
            assert self.w_zp_name
            packed_zp = torch.nn.Parameter(
                torch.ones(group_num, zp_output_size, dtype=torch.int32) * -2004318072,
                requires_grad=False,
            )
            setattr(layer, self.w_zp_name, packed_zp)

        # FIXME: some bugs in convert_weight_packed_scale_zp with GPTQ format,
        # repack to AWQ weight
        weight = unpack_quantized_values_into_int32(
            packed_weight, self.config.weight_type, 0
        )
        input_size, output_size = weight.size()
        weight = weight.view(input_size, output_size // 8, 8)
        weight = weight[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(input_size, output_size)
        weight = pack_quantized_values_into_int32(
            weight, self.config.weight_type, 1
        ).contiguous()

        zp = unpack_quantized_values_into_int32(packed_zp, self.config.weight_type, 1)
        zp = zp.view(group_num, output_size // 8, 8)
        zp = zp[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(group_num, output_size)
        zp = pack_quantized_values_into_int32(
            zp, self.config.weight_type, 1
        ).contiguous()

        blocked_w, blocked_zp, blocked_s = ops.convert_weight_packed_scale_zp(
            weight,
            zp,
            scales.data,
            ops.CPUQuantAlgo.AWQ,
        )

        if layer.bias is not None:
            layer.bias.data = layer.bias.float()

        packed_weight.data = blocked_w
        scales.data = blocked_s
        packed_zp.data = blocked_zp

    def process_weights_after_loading(self, layer: torch.nn.Module):
        if (not self.config.zero_points) and (self.w_zp_name is not None):
            setattr(layer, self.w_zp_name, None)

        if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
            setattr(layer, self.w_gidx_name, None)

        weights = getattr(layer, self.w_q_name)
        # Require GPTQ pack format
        assert weights.input_dim == weights.packed_dim

        # Weights in CT format is [output_size, input_size]
        if weights.input_dim == 1:
            weights.data = weights.t()

        # Scales in CT format is [output_size, group_num]
        scales = getattr(layer, self.w_s_name)
        if scales.output_dim == 0:
            scales.data = scales.t().contiguous()

        # Zero points in CT format is [output_size, group_num]
        # Zero points in awq_marlin format is [output_size, group_num]
        if self.config.zero_points:
            assert self.w_zp_name
            zp = getattr(layer, self.w_zp_name)
            if zp.output_dim == 0:
                zp.data = zp.t().contiguous()

        layer.use_w4a8 = (
            envs.VLLM_CPU_INT4_W4A8
            and not self.config.has_g_idx
            and self.config.act_type == torch.bfloat16
            and torch.cpu._is_amx_tile_supported()
        )
        # layer.use_w4a8 = False
        # AWQ format will be converted to GPTQ format in `AWQMarlinLinearMethod`
        if layer.use_w4a8:
            self._process_gptq_weights_w4a8(layer)
        else:
            self._process_gptq_weights_w4a16(layer)

    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
        if layer.use_w4a8:
            x = ops.int4_scaled_mm_cpu(
                x=x,
                w=w_q,
                w_zeros=w_zp,
                w_scales=w_s,
                bias=bias,
            )
        else:
            x = ops.cpu_gemm_wna16(
                input=x,
                q_weight=w_q,
                scales=w_s,
                zeros=w_zp,
                g_idx=w_gidx,
                bias=bias,
                pack_factor=8,  # 32 // 4
                isa_hint=layer.isa_hint,
            )
        return x


def _get_isa_hint(dtype: torch.dtype) -> str:
    supports_amx = torch.cpu._is_amx_tile_supported()
    if supports_amx and dtype in (torch.bfloat16,):
        return "amx"
    else:
        return "vec"
