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_optimized(
    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,
    NUM_STAGES: tl.constexpr,
):
    # AMD 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)
    
    # スーパーグループ化による効率向上
    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)
    offs_k = tl.arange(0, BLOCK_SIZE_K)
    
    # マスク事前計算
    m_mask = offs_m < M
    n_mask = offs_n < N
    
    # スケール事前計算（固定ブロックサイズ利用）
    scale_m_idx = offs_m
    scale_n_idx = offs_n // 128
    
    # 高精度アキュムレータ
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    
    # パイプライン化されたK次元ループ
    k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
    
    for k_idx in range(k_tiles):
        k_start = k_idx * BLOCK_SIZE_K
        k_offs = k_start + offs_k
        k_mask = k_offs < K
        
        # A行列プリフェッチ（M x K, column-major）
        a_ptrs = a_ptr + offs_m[:, None] * stride_am + k_offs[None, :] * stride_ak
        a_mask = m_mask[:, None] & k_mask[None, :]
        a_block = tl.load(a_ptrs, mask=a_mask, other=0.0)
        
        # B行列プリフェッチ（N x K, column-major）
        b_ptrs = b_ptr + offs_n[:, None] * stride_bn + k_offs[None, :] * stride_bk
        b_mask = n_mask[:, None] & k_mask[None, :]
        b_block = tl.load(b_ptrs, mask=b_mask, other=0.0)
        
        # スケールファクター効率的ロード
        scale_k_idx = k_start // 128
        scale_k_valid = scale_k_idx < tl.cdiv(K, 128)
        
        # A行列スケール（M x K//128）
        a_scale_ptrs = a_scale_ptr + scale_m_idx * stride_asm + scale_k_idx * stride_ask
        a_scale_mask = m_mask & scale_k_valid
        a_scale_vec = tl.load(a_scale_ptrs, mask=a_scale_mask, other=1.0)
        
        # B行列スケール（N//128 x K//128）
        b_scale_ptrs = b_scale_ptr + scale_n_idx * stride_bsn + scale_k_idx * stride_bsk
        b_scale_mask = (scale_n_idx < tl.cdiv(N, 128)) & scale_k_valid
        b_scale_vec = tl.load(b_scale_ptrs, mask=b_scale_mask, other=1.0)
        
        # 型変換とスケーリングの最適化
        a_fp32 = a_block.to(tl.float32)
        b_fp32 = b_block.to(tl.float32)
        
        # ブロードキャスト最適化
        a_scaled = a_fp32 * a_scale_vec[:, None]
        b_scaled = b_fp32 * b_scale_vec[:, None]
        
        # AMD最適化GEMM（Matrix Core利用）
        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 M * N * K < 1_000_000:  # 小サイズ
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 32, 64, 64
        GROUP_SIZE_M = 4
        NUM_STAGES = 2
    elif M * N * K < 10_000_000:  # 中サイズ
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 64
        GROUP_SIZE_M = 8
        NUM_STAGES = 3
    elif K > 4096:  # 大K次元
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 64, 128
        GROUP_SIZE_M = 8
        NUM_STAGES = 4
    else:  # 大サイズ一般
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 128, 64, 64
        GROUP_SIZE_M = 8
        NUM_STAGES = 3
    
    # 動的グリッド最適化
    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),
    )
    
    # 最適化されたカーネル起動
    fp8_gemm_kernel_optimized[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_STAGES=NUM_STAGES,
        num_warps=8,
        num_ctas=1,
    )
    
    return c

python eval.py test /tmp/tmpuh4ss81w