import torch
import triton
import triton.language as tl
from task import input_t, output_t
import os

os.environ["TRITON_LOG_LEVEL"] = "ERROR"
os.environ['PYTORCH_ROCM_ARCH'] = 'gfx942'

@triton.jit
def fp8_gemm_kernel_mi300x(
    a_ptr, b_ptr, a_scale_ptr, b_scale_ptr, c_ptr,
    M, N, K,
    stride_am, stride_ak,
    stride_bn, stride_bk,
    stride_asm, stride_ask,
    stride_bsn, stride_bsk,
    stride_cm, stride_cn,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    # MI300X最適化：効率的なプログラムID計算
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    
    # AMD Infinity Fabric最適化：8の倍数でグループ化
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + (pid % group_size_m)
    pid_n = (pid % num_pid_in_group) // group_size_m

    # ブロックオフセット計算（最適化）
    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    
    # FP32アキュムレータ（精度維持）
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    # K次元ループ（パイプライン最適化）
    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        k_start = k * BLOCK_SIZE_K
        offs_k = k_start + tl.arange(0, BLOCK_SIZE_K)
        
        # マスク計算（効率化）
        k_mask = offs_k < K
        a_mask = (offs_m[:, None] < M) & k_mask[None, :]
        b_mask = (offs_n[:, None] < N) & k_mask[None, :]
        
        # メモリアクセス最適化（coalesced access）
        a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
        b_ptrs = b_ptr + offs_n[:, None] * stride_bn + offs_k[None, :] * stride_bk
        
        # データロード（ベクトル化）
        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
        
        # スケールファクター計算（最適化）
        scale_k_block = k_start // 128
        
        # A行列スケール（効率的アクセス）
        a_scale_ptrs = a_scale_ptr + offs_m * stride_asm + scale_k_block * stride_ask
        a_scale_mask = (offs_m < M) & (scale_k_block < tl.cdiv(K, 128))
        a_scale = tl.load(a_scale_ptrs, mask=a_scale_mask, other=1.0)
        
        # B行列スケール（ブロックワイズ）
        scale_n_block = offs_n // 128
        b_scale_ptrs = b_scale_ptr + scale_n_block * stride_bsn + scale_k_block * stride_bsk
        b_scale_mask = (scale_n_block < tl.cdiv(N, 128)) & (scale_k_block < tl.cdiv(K, 128))
        b_scale = tl.load(b_scale_ptrs, mask=b_scale_mask, other=1.0)
        
        # FP8→FP32変換とスケーリング（融合）
        a_fp32 = a.to(tl.float32) * a_scale[:, None]
        b_fp32 = b.to(tl.float32) * b_scale[:, None]
        
        # MFMA最適化（MI300X特化）
        accumulator = tl.dot(a_fp32, tl.trans(b_fp32), accumulator, allow_tf32=False)
    
    # BF16変換（最終出力）
    c_result = accumulator.to(tl.bfloat16)
    
    # 出力ストア（効率化）
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + offs_cm[:, None] * stride_cm + offs_cn[None, :] * stride_cn
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    
    tl.store(c_ptrs, c_result, mask=c_mask)

@torch.compile(mode="max-autotune")
def custom_kernel(data: input_t) -> output_t:
    """
    ROCmドキュメント準拠MI300X最適化FP8 GEMM
    """
    a, b, a_scale, b_scale, c = data
    
    M, K = a.shape
    N, _ = b.shape
    
    # ストライド取得
    stride_am, stride_ak = a.stride()
    stride_bn, stride_bk = b.stride()
    stride_asm, stride_ask = a_scale.stride()
    stride_bsn, stride_bsk = b_scale.stride()
    stride_cm, stride_cn = c.stride()
    
    # MI300X最適化パラメータ（ROCmドキュメント準拠）
    # 最小1024ワークグループを目標
    total_elements = M * N
    
    if total_elements < 1_000_000:  # 小問題
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 32, 32, 64
        GROUP_SIZE_M = 8  # Infinity Fabric最適化
        num_warps = 4
        num_stages = 2  # 単一GEMM用
    elif total_elements < 10_000_000:  # 中問題
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 64, 64
        GROUP_SIZE_M = 8
        num_warps = 4
        num_stages = 2
    else:  # 大問題
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 64
        GROUP_SIZE_M = 8
        num_warps = 8
        num_stages = 2
    
    # ワークグループ数の計算（最適化）
    num_pid_m = triton.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = triton.cdiv(N, BLOCK_SIZE_N)
    total_workgroups = num_pid_m * num_pid_n
    
    # MI300X目標：最小1024ワークグループ
    if total_workgroups < 1024 and BLOCK_SIZE_M > 32:
        BLOCK_SIZE_M = max(32, BLOCK_SIZE_M // 2)
        BLOCK_SIZE_N = max(32, BLOCK_SIZE_N // 2)
    
    # グリッド計算
    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),
    )
    
    # MI300X最適化カーネル起動
    fp8_gemm_kernel_mi300x[grid](
        a, b, a_scale, b_scale, c,
        M, N, K,
        stride_am, stride_ak,
        stride_bn, stride_bk,
        stride_asm, stride_ask,
        stride_bsn, stride_bsk,
        stride_cm, stride_cn,
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
        BLOCK_SIZE_K=BLOCK_SIZE_K,
        GROUP_SIZE_M=GROUP_SIZE_M,
        num_warps=num_warps,
        num_stages=num_stages,
        matrix_instr_nonkdim=16,  # MFMA 16x16（MI300X最適）
        waves_per_eu=2,  # 占有率最適化
    )
    
    return c

python eval.py leaderboard /tmp/tmpxuym9k9l
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
skipping cudagraphs due to mutated inputs (1 instances)
