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(
    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,
):
    # プログラム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)
    
    # より効率的なグループ化
    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)
        
        # より正確なマスク計算
        a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K)
        b_mask = (offs_n[:, None] < N) & (offs_k[None, :] < K)
        
        # A行列データロード（M x K）
        a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
        a = tl.load(a_ptrs, mask=a_mask, other=0.0)
        
        # B行列データロード（N x K）
        b_ptrs = b_ptr + offs_n[:, None] * stride_bn + offs_k[None, :] * stride_bk
        b = tl.load(b_ptrs, mask=b_mask, other=0.0)
        
        # スケールファクター計算の修正
        scale_k_block = k_start // 128
        scale_k_mask = scale_k_block < tl.cdiv(K, 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_mask
        a_scale = tl.load(a_scale_ptrs, mask=a_scale_mask, other=1.0)
        
        # B行列スケールファクター（128要素ブロック単位）
        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_mask
        b_scale = tl.load(b_scale_ptrs, mask=b_scale_mask, other=1.0)
        
        # FP8からFP32への変換とスケーリング
        a_fp32 = a.to(tl.float32)
        b_fp32 = b.to(tl.float32)
        
        # スケールファクターの正確な適用
        a_scaled = a_fp32 * a_scale[:, None]
        b_scaled = b_fp32 * b_scale[:, None]
        
        # 高精度行列乗算
        accumulator = tl.dot(a_scaled, tl.trans(b_scaled), accumulator, allow_tf32=False)
    
    # 最終結果の変換
    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:
    """
    修正されたAMD 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用最適化パラメータ
    if K <= 512:
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 64, 64
        GROUP_SIZE_M = 4
    elif K <= 2048:
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 128
        GROUP_SIZE_M = 8
    else:
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 128, 128
        GROUP_SIZE_M = 8
    
    # グリッド計算
    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),
    )
    
    # カーネル起動
    fp8_gemm_kernel[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,
    )
    
    return c

python eval.py test /tmp/tmpyeaxun1n