# submission.py
import torch
from task import input_t, output_t
from aiter.ops import gemm_op_a8w8 as aiter_ck_gemm_ops


# ------------------------------------------------------------- #
# Kernel IDs for the nine benchmark shapes (split‑K = 0)
TUNED_KERNELS: dict[tuple[int, int, int], int] = {
    (1024, 1536, 7168): 18,
    (1024, 3072, 1536): 18,
    (1024, 576, 7168): 18,
    (1024, 7168, 256): 16,
    (1024, 7168, 2048): 2,
    (1024, 4608, 7168): 2,
    (1024, 7168, 2304): 2,
    (1024, 512, 7168): 18,
    (1024, 4096, 512): 1,
}
# ------------------------------------------------------------- #


# Helper: keep exactly the same name and behaviour you had
def ensure_row_major(t: torch.Tensor) -> torch.Tensor:
    if not t.is_contiguous(memory_format=torch.contiguous_format):
        return t.clone(memory_format=torch.contiguous_format)
    return t


def custom_kernel(data: input_t) -> output_t:
    A_cm, B_cm, alpha_cm, beta_cm, C_cm_orig = data

    # Convert inputs (unchanged from your original code)
    A_rm = ensure_row_major(A_cm).to(torch.float8_e4m3fnuz)
    B_rm = ensure_row_major(B_cm).to(torch.float8_e4m3fnuz)
    alpha_rm = ensure_row_major(alpha_cm).to(torch.float32)
    beta_rm = ensure_row_major(beta_cm).to(torch.float32)

    # Decide whether this (M,N,K) is in the tuned list
    m, k = A_cm.shape
    n = B_cm.shape[0]
    key = (m, n, k)

    if key in TUNED_KERNELS:
        # Call the *tune* overload → extra kernelId & splitK args
        aiter_ck_gemm_ops.gemm_a8w8_blockscale_tune(
            A_rm,
            B_rm,
            alpha_rm,
            beta_rm,
            C_cm_orig,
            TUNED_KERNELS[key],  # kernelId
            0,  # splitK fixed to 0
        )
    else:
        # Fall back to the generic dispatcher
        aiter_ck_gemm_ops.gemm_a8w8_blockscale(A_rm, B_rm, alpha_rm, beta_rm, C_cm_orig)

    return C_cm_orig


# k: 7168; m: 1024; n: 1536; seed: 8135
#  ⏱ 176 ± 0.2 µs
#  ⚡ 175 µs 🐌 178 µs

# k: 1536; m: 1024; n: 3072; seed: 6251
#  ⏱ 88.2 ± 0.10 µs
#  ⚡ 86.6 µs 🐌 93.6 µs

# k: 7168; m: 1024; n: 576; seed: 12346
#  ⏱ 111 ± 0.1 µs
#  ⚡ 110 µs 🐌 117 µs

# k: 256; m: 1024; n: 7168; seed: 5364
#  ⏱ 76.4 ± 0.19 µs
#  ⚡ 73.4 µs 🐌 86.3 µs

# k: 2048; m: 1024; n: 7168; seed: 6132
#  ⏱ 172 ± 0.1 µs
#  ⚡ 172 µs 🐌 173 µs

# k: 7168; m: 1024; n: 4608; seed: 7531
#  ⏱ 378 ± 0.7 µs
#  ⚡ 364 µs 🐌 393 µs

# k: 2304; m: 1024; n: 7168; seed: 12345
#  ⏱ 193 ± 0.2 µs
#  ⚡ 189 µs 🐌 198 µs

# k: 7168; m: 1024; n: 512; seed: 6563
#  ⏱ 111 ± 0.5 µs
#  ⚡ 108 µs 🐌 142 µs

# k: 512; m: 1024; n: 4096; seed: 17512
#  ⏱ 76.2 ± 0.17 µs
#  ⚡ 73.3 µs 🐌 87.5 µs

# k: 7168; m: 6144; n: 1536; seed: 6543
#  ⏱ 711 ± 2.4 µs
#  ⚡ 669 µs 🐌 771 µs

# k: 1536; m: 6144; n: 3072; seed: 234
#  ⏱ 233 ± 0.2 µs
#  ⚡ 230 µs 🐌 239 µs

# k: 7168; m: 6144; n: 576; seed: 9863
#  ⏱ 512 ± 2.2 µs
#  ⚡ 476 µs 🐌 550 µs

# k: 256; m: 6144; n: 7168; seed: 764243
#  ⏱ 131 ± 0.2 µs
#  ⚡ 128 µs 🐌 138 µs

# k: 2048; m: 6144; n: 7168; seed: 76547
#  ⏱ 510 ± 1.1 µs
#  ⚡ 496 µs 🐌 534 µs

# k: 7168; m: 6144; n: 4608; seed: 65436
#  ⏱ 1318 ± 3.2 µs
#  ⚡ 1251 µs 🐌 1376 µs

# k: 2304; m: 6144; n: 7168; seed: 452345
#  ⏱ 597 ± 0.5 µs
#  ⚡ 595 µs 🐌 598 µs

# k: 7168; m: 6144; n: 512; seed: 12341
#  ⏱ 507 ± 2.4 µs
#  ⚡ 469 µs 🐌 559 µs

# k: 512; m: 6144; n: 4096; seed: 45245
#  ⏱ 131 ± 0.2 µs
#  ⚡ 128 µs 🐌 141 µs

# Successful on MI300: 233 µs

python eval.py benchmark /tmp/tmpiz2hcu_f
[aiter] start build [module_gemm_a8w8_blockscale_tune] under /home/runner/aiter/aiter/jit/build/module_gemm_a8w8_blockscale_tune
[92mSuccessfully preprocessed all matching files.[0m
[aiter] finish build [module_gemm_a8w8_blockscale_tune], cost 93.17133053s
[aiter] type hints mismatch, override to --> gemm_a8w8_blockscale_tune(XQ: torch.Tensor, WQ: torch.Tensor, x_scale: torch.Tensor, w_scale: torch.Tensor, Out: torch.Tensor, kernelId: int = 0, splitK: int = 0) -> torch.Tensor
[aiter] type hints mismatch, override to --> gemm_a8w8_blockscale(XQ: torch.Tensor, WQ: torch.Tensor, x_scale: torch.Tensor, w_scale: torch.Tensor, Out: torch.Tensor) -> torch.Tensor
