import torch

from kernels.benchmark import Benchmark


def _quantize_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return torch.clamp(x.float() / scale.float(), -448.0, 448.0).to(torch.float8_e4m3fn)


def _dequant_fp8(x: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return x.float() * scale.float()


class FP8SwiGLUMlpBenchmark(Benchmark):
    seed = 17

    def _setup_shape(self, M: int, K: int, H: int, N: int) -> None:
        self.M, self.K, self.H, self.N = M, K, H, N
        self.x_scale = torch.tensor([0.05], device=self.device, dtype=torch.float32)
        self.gate_up_scale = torch.tensor([0.04], device=self.device, dtype=torch.float32)
        self.hidden_scale = torch.tensor([0.25], device=self.device, dtype=torch.float32)
        self.down_scale = torch.tensor([0.04], device=self.device, dtype=torch.float32)
        self.x = _quantize_fp8(
            torch.randn((M, K), device=self.device, dtype=torch.bfloat16),
            self.x_scale,
        )
        self.gate_up_w = _quantize_fp8(
            torch.randn((2 * H, K), device=self.device, dtype=torch.bfloat16),
            self.gate_up_scale,
        )
        self.down_w = _quantize_fp8(
            torch.randn((N, H), device=self.device, dtype=torch.bfloat16),
            self.down_scale,
        )
        self.gate_up_bf16 = torch.empty((M, 2 * H), device=self.device, dtype=torch.bfloat16)
        self.hidden_fp8 = torch.empty((M, H), device=self.device, dtype=torch.float8_e4m3fn)
        self.out = torch.empty((M, N), device=self.device, dtype=torch.bfloat16)

    def _reference(self) -> torch.Tensor:
        gate_up = (
            _dequant_fp8(self.x, self.x_scale)
            @ _dequant_fp8(self.gate_up_w, self.gate_up_scale).T
        ).to(torch.bfloat16)
        gate, up = gate_up.float().chunk(2, dim=1)
        hidden_fp8 = _quantize_fp8(torch.nn.functional.silu(gate) * up, self.hidden_scale)
        return (
            _dequant_fp8(hidden_fp8, self.hidden_scale)
            @ _dequant_fp8(self.down_w, self.down_scale).T
        ).to(torch.bfloat16)

    def setup_smoke_mlp(self) -> None:
        self._setup_shape(10, 1024, 4096, 1024)

    def benchmark_smoke_mlp(self) -> None:
        self.kernel.fp8_swiglu_mlp_bf16(
            self.x,
            self.gate_up_w,
            self.down_w,
            self.x_scale,
            self.gate_up_scale,
            self.hidden_scale,
            self.down_scale,
            gate_up_bf16=self.gate_up_bf16,
            hidden_fp8=self.hidden_fp8,
            out=self.out,
        )

    def verify_smoke_mlp(self) -> torch.Tensor:
        return self._reference()