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

# This registers op implementations
import vllm.kernels  # noqa: F401
from tests.ir.ir_test_utils import (
    COMMON_HIDDEN_SIZES,
    NUM_TOKENS,
    assert_close,
    clone_args,
    supported_providers,
)
from tests.kernels.allclose_default import get_default_rtol
from vllm import ir
from vllm.platforms import current_platform

rms_norm_native = ir.ops.rms_norm.impls["native"].impl_fn


@pytest.mark.skipif(
    not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
    reason="Currently only kernels on CUDA, ROCm and XPU",
)
def test_rms_norm_registration():
    expected = {
        "native": True,
        "vllm_c": current_platform.is_cuda_alike(),
        "aiter": current_platform.is_rocm(),
        "oink": current_platform.has_device_capability(100)
        and hasattr(torch.ops, "oink")
        and hasattr(torch.ops.oink, "rmsnorm"),
        "xpu_kernels": current_platform.is_xpu(),
    }

    actual = {
        provider: impl.supported for provider, impl in ir.ops.rms_norm.impls.items()
    }

    assert actual == expected


@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("n_tokens", NUM_TOKENS)
@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
@pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
@pytest.mark.skipif(
    not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
    reason="Currently only kernels on CUDA, ROCm and XPU",
)
class TestRMSNorm:
    @classmethod
    def setup_class(cls, **kwargs):
        torch.set_default_device(current_platform.device_type)

    def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
        x, weight, epsilon = ir.ops.rms_norm.generate_inputs(
            num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
        )
        out = rms_norm_native(x, weight, epsilon=epsilon)

        # Check shape, dtype, device
        assert out.shape == x.shape
        assert out.dtype == x.dtype
        assert out.device == x.device

        # Check the scaling property of rms norm
        out2 = rms_norm_native(x * 2.0, weight, epsilon=epsilon)
        torch.testing.assert_close(out2, out, rtol=get_default_rtol(out), atol=1e-3)

        # Mean square should be approximately 1 (ignoring epsilon and weight scaling)
        combined_norm = out.float() / weight.float()
        variance = combined_norm.pow(2).mean(dim=-1)
        # After RMS normalization, variance should be close to 1
        torch.testing.assert_close(
            variance, torch.ones_like(variance), rtol=1e-2, atol=1e-2
        )

        # Check behavior with and without weight
        weight1 = torch.ones_like(weight)
        out3 = rms_norm_native(x, weight1, epsilon=epsilon)
        out4 = rms_norm_native(x, None, epsilon=epsilon)
        torch.testing.assert_close(out3, out4)

    @pytest.mark.parametrize("provider", supported_providers(ir.ops.rms_norm))
    def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
        impl = ir.ops.rms_norm.impls[provider]
        x, weight, eps = ir.ops.rms_norm.generate_inputs(
            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
        )
        args = (x, weight, eps)

        if not impl.supports_args(*args):
            pytest.skip(f"{provider} does not support args")

        ref_output = rms_norm_native(*clone_args(args))
        output = impl.impl_fn(*clone_args(args))
        assert_close(ir.ops.rms_norm, output, ref_output)

        # check that dispatched call matches direct call
        with ir.ops.rms_norm.set_priority([provider, "native"]):
            out_dispatched = ir.ops.rms_norm(*args)
        out_direct = impl.impl_fn(*args)
        torch.testing.assert_close(out_dispatched, out_direct, rtol=0.0, atol=0.0)

        # none of these support variance_size override
        assert not impl.supports_args(x, weight, eps, 4)
        assert not impl.supports_args(x, weight, eps, variance_size=4)

        # test weight=None behavior
        out_no_weight = impl.impl_fn(x, None, eps)
        out_unit_weight = impl.impl_fn(x, torch.ones_like(weight), eps)
        assert_close(ir.ops.rms_norm, out_no_weight, out_unit_weight)

    @pytest.mark.parametrize("provider", ["vllm_c", "aiter", "xpu_kernels", "native"])
    def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
        if not ir.ops.rms_norm.impls[provider].supported:
            pytest.skip(f"{provider} impl not supported on this platform")

        args = ir.ops.rms_norm.generate_inputs(
            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
        )

        # When checking the torch op, we have to set priority and use dispatch
        with ir.ops.rms_norm.set_priority([provider, "native"]):
            torch.library.opcheck(torch.ops.vllm_ir.rms_norm, args)


@pytest.mark.skipif(
    not current_platform.is_rocm(),
    reason="aiter is only supported on ROCm",
)
def test_aiter_rejects_unsupported_dtypes():
    torch.set_default_device(current_platform.device_type)
    impl = ir.ops.rms_norm.impls["aiter"]
    for dtype in [torch.float32, torch.float64]:
        args = ir.ops.rms_norm.generate_inputs(
            num_tokens=8, hidden_size=4096, dtype=dtype, epsilon=1e-5
        )
        assert not impl.supports_args(*args), f"aiter should reject dtype={dtype}"


@pytest.mark.skipif(
    not current_platform.is_rocm(),
    reason="ROCm vllm_c RMSNorm needs explicit ND input handling",
)
def test_vllm_c_rms_norm_accepts_nd_input():
    torch.set_default_device(current_platform.device_type)
    impl = ir.ops.rms_norm.impls["vllm_c"]
    if not impl.supported:
        pytest.skip("vllm_c impl not supported on this platform")

    base = torch.randn(3, 8, 192, dtype=torch.float16)
    x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
    assert not x.is_contiguous()
    weight = torch.randn(16, dtype=torch.float16)
    epsilon = 1e-5

    output = impl.impl_fn(x, weight, epsilon)
    ref_output = rms_norm_native(x, weight, epsilon)

    assert output.shape == x.shape
    assert_close(ir.ops.rms_norm, output, ref_output)


fused_add_rms_norm_native = ir.ops.fused_add_rms_norm.impls["native"].impl_fn


@pytest.mark.skipif(
    not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
    reason="Currently only kernels on CUDA, ROCm and XPU",
)
def test_fused_add_rms_norm_registration():
    expected = {
        "native": True,
        "vllm_c": current_platform.is_cuda_alike(),
        "aiter": current_platform.is_rocm(),
        "oink": current_platform.has_device_capability(100)
        and hasattr(torch.ops, "oink")
        and hasattr(torch.ops.oink, "fused_add_rms_norm"),
        "xpu_kernels": current_platform.is_xpu(),
    }

    actual = {
        provider: impl.supported
        for provider, impl in ir.ops.fused_add_rms_norm.impls.items()
    }

    assert actual == expected


@pytest.mark.skipif(
    not current_platform.is_rocm(),
    reason="ROCm vllm_c fused_add_rms_norm needs explicit ND input handling",
)
def test_vllm_c_fused_add_rms_norm_accepts_nd_input():
    torch.set_default_device(current_platform.device_type)
    impl = ir.ops.fused_add_rms_norm.impls["vllm_c"]
    if not impl.supported:
        pytest.skip("vllm_c impl not supported on this platform")

    base = torch.randn(3, 8, 192, dtype=torch.float16)
    residual_base = torch.randn(3, 8, 192, dtype=torch.float16)
    x = base.split(64, dim=-1)[0].view(3, 8, 4, 16)
    x_residual = residual_base.split(64, dim=-1)[0].view(3, 8, 4, 16)
    assert not x.is_contiguous()
    assert not x_residual.is_contiguous()
    weight = torch.randn(16, dtype=torch.float16)
    epsilon = 1e-5

    output, residual = impl.impl_fn(x.clone(), x_residual.clone(), weight, epsilon)
    ref_output, ref_residual = fused_add_rms_norm_native(x, x_residual, weight, epsilon)

    assert output.shape == x.shape
    assert residual.shape == x_residual.shape
    assert_close(ir.ops.fused_add_rms_norm, output, ref_output)
    assert_close(ir.ops.fused_add_rms_norm, residual, ref_residual)


@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("n_tokens", NUM_TOKENS)
@pytest.mark.parametrize("hidden_size", COMMON_HIDDEN_SIZES)
@pytest.mark.parametrize("epsilon", [1e-6, 1e-5])
@pytest.mark.skipif(
    not current_platform.is_cuda_alike() and not current_platform.is_xpu(),
    reason="Currently only kernels on CUDA, ROCm and XPU",
)
class TestFusedAddRMSNorm:
    @classmethod
    def setup_class(cls, **kwargs):
        torch.set_default_device(current_platform.device_type)

    def test_native_semantics(self, dtype, n_tokens, hidden_size, epsilon):
        x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
            num_tokens=4, hidden_size=8, dtype=dtype, epsilon=epsilon
        )
        out, residual_out = fused_add_rms_norm_native(x, x_residual, weight, eps)

        # Check shape, dtype, device
        assert out.shape == x.shape
        assert out.dtype == x.dtype
        assert out.device == x.device
        assert residual_out.shape == x_residual.shape
        assert residual_out.dtype == x_residual.dtype
        assert residual_out.device == x_residual.device

        # Check that residual_out = x + x_residual
        expected_residual = (x.float() + x_residual.float()).to(dtype)
        torch.testing.assert_close(
            residual_out, expected_residual, rtol=1e-3, atol=1e-3
        )

        # Verify that the output is RMS normalized version of (x + x_residual)
        expected_out = rms_norm_native(expected_residual, weight, epsilon)
        assert_close(
            ir.ops.fused_add_rms_norm,
            (out, residual_out),
            (expected_out, expected_residual),
        )

        # Check the scaling property of rms norm
        out1, _ = fused_add_rms_norm_native(
            x, torch.zeros_like(x), weight, epsilon=epsilon
        )
        out2, _ = fused_add_rms_norm_native(
            x * 2.0, torch.zeros_like(x), weight, epsilon=epsilon
        )
        torch.testing.assert_close(out2, out1, rtol=get_default_rtol(out), atol=1e-3)

        # Check behavior with and without weight
        weight1 = torch.ones_like(weight)
        out3, _ = fused_add_rms_norm_native(x, x_residual, weight1, eps)
        out4, _ = fused_add_rms_norm_native(x, x_residual, None, eps)
        torch.testing.assert_close(out3, out4)

    @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
    def test_impls(self, dtype, n_tokens, hidden_size, epsilon, provider):
        impl = ir.ops.fused_add_rms_norm.impls[provider]
        x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
        )
        args = (x, x_residual, weight, eps, None)

        if not impl.supports_args(*args):
            pytest.skip(f"{provider} does not support args")

        ref_output, ref_residual = fused_add_rms_norm_native(*clone_args(args))
        output, residual = impl.impl_fn(*clone_args(args))
        assert_close(ir.ops.fused_add_rms_norm, output, ref_output)
        assert_close(ir.ops.fused_add_rms_norm, residual, ref_residual)

        # check that dispatched call matches direct call
        with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
            out_dispatched, residual_dispatched = ir.ops.fused_add_rms_norm(*args[:4])
        out_direct, residual_direct = impl.impl_fn(*clone_args(args))
        torch.testing.assert_close(out_dispatched, out_direct, rtol=0.0, atol=0.0)
        torch.testing.assert_close(
            residual_dispatched, residual_direct, rtol=0.0, atol=0.0
        )

        # none of these support variance_size override
        assert not impl.supports_args(x, x_residual, weight, epsilon, 4)
        assert not impl.supports_args(x, x_residual, weight, epsilon, variance_size=4)

        # test weight=None behavior
        out_no_weight, residual_no_weight = impl.impl_fn(
            x.clone(), x_residual.clone(), None, epsilon
        )
        out_unit_weight, residual_unit_weight = impl.impl_fn(
            x.clone(), x_residual.clone(), torch.ones_like(weight), epsilon
        )
        assert_close(ir.ops.fused_add_rms_norm, out_no_weight, out_unit_weight)
        assert_close(
            ir.ops.fused_add_rms_norm, residual_no_weight, residual_unit_weight
        )

    @pytest.mark.parametrize("provider", ["vllm_c"])
    def test_inplace_semantics(self, dtype, n_tokens, hidden_size, epsilon, provider):
        """Test that inplace implementations reuse inputs,
        for maybe_inplace overload but not for default overload."""
        impl = ir.ops.fused_add_rms_norm.impls[provider]
        if not impl.supported:
            pytest.skip(f"{provider} impl not supported on this platform")

        x, x_residual, weight, eps = ir.ops.fused_add_rms_norm.generate_inputs(
            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
        )

        # Test default overload - should NOT modify inputs even with inplace impl
        x_default = x.clone()
        x_residual_default = x_residual.clone()
        x_default_ptr = x_default.data_ptr()
        x_residual_default_ptr = x_residual_default.data_ptr()

        with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
            out_default, residual_default = ir.ops.fused_add_rms_norm(
                x_default, x_residual_default, weight, eps
            )

        # Default should NOT be inplace (even with inplace implementation)
        assert out_default.data_ptr() != x_default_ptr
        assert residual_default.data_ptr() != x_residual_default_ptr
        torch.testing.assert_close(x, x_default, rtol=0.0, atol=0.0)
        torch.testing.assert_close(x_residual, x_residual_default, rtol=0.0, atol=0.0)

        # Test maybe_inplace overload - should modify inputs with inplace impl
        x_inplace = x.clone()
        x_residual_inplace = x_residual.clone()
        x_inplace_ptr = x_inplace.data_ptr()
        x_residual_inplace_ptr = x_residual_inplace.data_ptr()

        with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
            out_inplace, residual_inplace = ir.ops.fused_add_rms_norm.maybe_inplace(
                x_inplace, x_residual_inplace, weight, eps
            )

        # maybe_inplace should be inplace
        assert out_inplace.data_ptr() == x_inplace_ptr
        assert residual_inplace.data_ptr() == x_residual_inplace_ptr

        # Both should produce same results
        torch.testing.assert_close(out_default, out_inplace, atol=0.0, rtol=0.0)
        torch.testing.assert_close(
            residual_default, residual_inplace, atol=0.0, rtol=0.0
        )

    @pytest.mark.parametrize("provider", supported_providers(ir.ops.fused_add_rms_norm))
    def test_torch_opcheck(self, dtype, n_tokens, hidden_size, epsilon, provider):
        args = ir.ops.fused_add_rms_norm.generate_inputs(
            num_tokens=n_tokens, hidden_size=hidden_size, dtype=dtype, epsilon=epsilon
        )
        args = args + (None,)  # Add variance_size parameter

        # When checking the torch op, we have to set priority and use dispatch
        with ir.ops.fused_add_rms_norm.set_priority([provider, "native"]):
            torch.library.opcheck(torch.ops.vllm_ir.fused_add_rms_norm.default, args)

            # Only test maybe_inplace with non-inplace implementations
            # Inplace implementations return aliases of inputs which is not allowed.
            # We break this invariant, but we also convert maybe_inplace to the default
            # overload during compilation, so maybe_inplace never reaches Inductor.
            if not ir.ops.fused_add_rms_norm.impls[provider].inplace:
                torch.library.opcheck(
                    torch.ops.vllm_ir.fused_add_rms_norm.maybe_inplace, args
                )