from enum import Enum
from itertools import product

import torch


f_ref = torch._scaled_grouped_mm
f = torch.compile(
    f_ref,
    options={
        "max_autotune": True,
        "max_autotune_gemm_backends": "TRITON",
    },
)


class MMType(Enum):
    MM_2D_2D = 1
    MM_2D_3D = 2
    MM_3D_2D = 3
    MM_3D_3D = 4


def generate_data(
    mm_type, group_size, M, N, K, device, dtype_AB, dtype_scale, dtype_offset, strided
):
    if mm_type == MMType.MM_2D_2D:
        A = torch.randn(M, K * (group_size + strided), device=device).to(dtype_AB)[
            :, : K * group_size
        ]
        B = torch.randn(N, K * (group_size + strided), device=device).to(dtype_AB)[
            :, : K * group_size
        ]
        A_scale = torch.rand(group_size * M, device=device, dtype=dtype_scale)
        B_scale = torch.rand(group_size * N, device=device, dtype=dtype_scale)
        offs = torch.arange(K, group_size * K + 1, K, device=device, dtype=dtype_offset)

    if mm_type == MMType.MM_2D_3D:
        A = torch.randn(M * group_size, K * (1 + strided), device=device).to(dtype_AB)[
            :, :K
        ]
        B = torch.randn(
            group_size * (1 + strided), N, K * (1 + strided), device=device
        ).to(dtype_AB)[:: (1 + strided), :, :K]
        A_scale = torch.rand(group_size * M, device=device, dtype=dtype_scale)
        B_scale = torch.rand(
            group_size, N * (1 + strided), device=device, dtype=dtype_scale
        )[:, :N]
        offs = torch.arange(M, group_size * M + 1, M, device=device, dtype=dtype_offset)

    if mm_type == MMType.MM_3D_2D:
        A = torch.randn(
            group_size * (1 + strided), M, K * (1 + strided), device=device
        ).to(dtype_AB)[:: (1 + strided), :, :K]
        B = torch.randn(N * group_size, K * (1 + strided), device=device).to(dtype_AB)[
            :, :K
        ]
        A_scale = torch.rand(
            group_size, M * (1 + strided), device=device, dtype=dtype_scale
        )[:, :M]
        B_scale = torch.rand(group_size * N, device=device, dtype=dtype_scale)
        offs = torch.arange(N, group_size * N + 1, N, device=device, dtype=dtype_offset)

    if mm_type == MMType.MM_3D_3D:
        A = torch.randn(
            group_size * (1 + strided), M, K * (1 + strided), device=device
        ).to(dtype_AB)[:: (1 + strided), :, :K]
        B = torch.randn(
            group_size * (1 + strided), N, K * (1 + strided), device=device
        ).to(dtype_AB)[:: (1 + strided), :, :K]
        A_scale = torch.rand(group_size, M * (1 + strided), device=device).to(
            dtype_scale
        )[:, :M]
        B_scale = torch.rand(group_size, N * (1 + strided), device=device).to(
            dtype_scale
        )[:, :N]
        offs = None

    if offs is not None:
        if offs[0] >= 32:
            offs[0] -= 16
            offs[2] += 16
        elif offs[0] >= 64:
            offs[0] -= 16
            offs[1] += 16
            offs[2] -= 32

    return A, B, A_scale, B_scale, offs


def validate():
    def validate_helper(
        mm_type,
        group_size,
        M,
        N,
        K,
        device,
        dtype_AB,
        dtype_scale,
        dtype_offset,
        dtype_C,
        use_fast_accum,
        strided,
        atol,
        rtol,
    ):
        torch._dynamo.reset()

        A, B, A_scale, B_scale, offs = generate_data(
            mm_type,
            group_size,
            M,
            N,
            K,
            device,
            dtype_AB,
            dtype_scale,
            dtype_offset,
            strided,
        )

        C_ref = f_ref(
            A,
            B.transpose(-2, -1),
            A_scale,
            B_scale,
            offs,
            out_dtype=dtype_C,
            use_fast_accum=use_fast_accum,
        )
        C = f(
            A,
            B.transpose(-2, -1),
            A_scale,
            B_scale,
            offs,
            out_dtype=dtype_C,
            use_fast_accum=use_fast_accum,
        )
        assert torch.allclose(C, C_ref, atol=atol, rtol=rtol)

    device = "cuda"
    group_size = 4
    M_range = [2**i for i in range(4, 6)]
    N_range = [2**i for i in range(5, 8)]
    K_range = [2**i for i in range(6, 9)]
    dtype_AB = torch.float8_e4m3fn
    dtype_scale = torch.float32
    dtype_offset = torch.int32
    dtype_C = torch.bfloat16
    use_fast_accum_range = [False, True]
    strided_range = [False, True]
    atol = 1e-2
    rtol = 1e-2
    for mm_type, M, N, K, use_fast_accum, strided in product(
        MMType, M_range, N_range, K_range, use_fast_accum_range, strided_range
    ):
        validate_helper(
            mm_type,
            group_size,
            M,
            N,
            K,
            device,
            dtype_AB,
            dtype_scale,
            dtype_offset,
            dtype_C,
            use_fast_accum,
            strided,
            atol,
            rtol,
        )


validate()