import torch
import torch.nn.functional as F

from kernels.benchmark import Benchmark


class ReluBenchmark(Benchmark):
    """Same workloads/names as upstream kernels-community/relu (fp32
    base/large) so results are directly comparable across the two repos."""

    seed: int = 42

    def setup(self):
        self.x = torch.randn(1024, 1024, device=self.device, dtype=torch.float32)
        self.out = torch.empty_like(self.x)

    def benchmark_base(self):
        self.out = self.kernel.relu(self.x)

    def verify_base(self) -> torch.Tensor:
        return F.relu(self.x)

    def setup_large(self):
        self.x = torch.randn(4096, 4096, device=self.device, dtype=torch.float32)
        self.out = torch.empty_like(self.x)

    def benchmark_large(self):
        self.out = self.kernel.relu(self.x)

    def verify_large(self) -> torch.Tensor:
        return F.relu(self.x)


class ReluDtypeBenchmark(Benchmark):
    """Dtypes this build ships beyond upstream's fp32-only: fp16, bf16, int8.
    base = 1024^2, large = 4096^2, matching ReluBenchmark's shapes.

    Uses the kernel's out= form with a preallocated output: the allocating
    form (`self.out = kernel.relu(x)`) keeps the previous output alive at
    allocation time, so the caching allocator alternates two output blocks
    and the effective footprint grows by a full buffer — at L2-boundary
    working sets that alone pushes the timed loop off the L2 cliff while
    the single-shot reference timing stays on it (torch.relu itself
    measures ~0.33x vs its own single-shot under the same loop). The out=
    form keeps one live output, matching the reference's footprint."""

    seed: int = 42

    def _float_setup(self, side, dtype):
        self.x = torch.randn(side, side, device=self.device, dtype=dtype)
        self.out = torch.empty_like(self.x)
        self.ref_out = torch.empty_like(self.x)

    def setup_fp16_base(self):
        self._float_setup(1024, torch.float16)

    def setup_fp16_large(self):
        self._float_setup(4096, torch.float16)

    def setup_bf16_base(self):
        self._float_setup(1024, torch.bfloat16)

    def setup_bf16_large(self):
        self._float_setup(4096, torch.bfloat16)

    def _int8_setup(self, side):
        self.x = torch.randint(
            -128, 128, (side, side), device=self.device, dtype=torch.int8
        )
        self.out = torch.empty_like(self.x)
        self.ref_out = torch.empty_like(self.x)

    def setup_int8_base(self):
        self._int8_setup(1024)

    def setup_int8_large(self):
        self._int8_setup(4096)

    def benchmark_fp16_base(self):
        self.kernel.relu(self.x, out=self.out)

    def benchmark_fp16_large(self):
        self.kernel.relu(self.x, out=self.out)

    def benchmark_bf16_base(self):
        self.kernel.relu(self.x, out=self.out)

    def benchmark_bf16_large(self):
        self.kernel.relu(self.x, out=self.out)

    def benchmark_int8_base(self):
        self.kernel.relu(self.x, out=self.out)

    def benchmark_int8_large(self):
        self.kernel.relu(self.x, out=self.out)

    def verify_fp16_base(self) -> torch.Tensor:
        torch.clamp(self.x, min=0, out=self.ref_out)
        return self.ref_out

    def verify_fp16_large(self) -> torch.Tensor:
        torch.clamp(self.x, min=0, out=self.ref_out)
        return self.ref_out

    def verify_bf16_base(self) -> torch.Tensor:
        torch.clamp(self.x, min=0, out=self.ref_out)
        return self.ref_out

    def verify_bf16_large(self) -> torch.Tensor:
        torch.clamp(self.x, min=0, out=self.ref_out)
        return self.ref_out

    def verify_int8_base(self) -> torch.Tensor:
        torch.clamp(self.x, min=0, out=self.ref_out)
        return self.ref_out

    def verify_int8_large(self) -> torch.Tensor:
        torch.clamp(self.x, min=0, out=self.ref_out)
        return self.ref_out