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)

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 benchmark /tmp/tmp_jf_jls9