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_rocwmma_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,
):
    # rocWMMA最適化：wavefront単位の効率的なプログラム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)
    
    # MFMA wavefront最適化グループ化（64ワープ = 1 CU）
    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

    # rocWMMA fragment形式のオフセット計算
    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)
    
    # MFMA最適化：FP32アキュムレータ（CDNA3推奨）
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    # rocWMMA style: 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)
        
        # Fragment単位のマスク計算
        k_valid = offs_k < K
        a_mask = (offs_m[:, None] < M) & k_valid[None, :]
        b_mask = (offs_n[:, None] < N) & k_valid[None, :]
        
        # MFMA最適化メモリアクセス（wavefront単位）
        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
        
        # rocWMMA推奨：ベクトル化ロード
        a_fragment = tl.load(a_ptrs, mask=a_mask, other=0.0)
        b_fragment = tl.load(b_ptrs, mask=b_mask, other=0.0)
        
        # FP8スケーリング最適化（rocWMMA準拠）
        scale_k_idx = k_start // 128
        
        # A行列スケール（MFMA fragment対応）
        a_scale_ptrs = a_scale_ptr + offs_m * stride_asm + scale_k_idx * stride_ask
        a_scale_mask = (offs_m < M) & (scale_k_idx < tl.cdiv(K, 128))
        a_scale_frag = tl.load(a_scale_ptrs, mask=a_scale_mask, other=1.0)
        
        # B行列スケール（wavefront分散）
        scale_n_idx = offs_n // 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_idx < tl.cdiv(K, 128))
        b_scale_frag = tl.load(b_scale_ptrs, mask=b_scale_mask, other=1.0)
        
        # rocWMMA推奨：FP8→FP32変換とスケーリング融合
        a_fp32 = a_fragment.to(tl.float32) * a_scale_frag[:, None]
        b_fp32 = b_fragment.to(tl.float32) * b_scale_frag[:, None]
        
        # MFMA最適化：AMD CDNA3のieee精度使用
        accumulator = tl.dot(a_fp32, tl.trans(b_fp32), accumulator, input_precision="ieee")
    
    # rocWMMA推奨：BF16最終出力
    c_result = accumulator.to(tl.bfloat16)
    
    # MFMA最適化ストア（wavefront効率）
    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(fullgraph=True,mode="max-autotune")
def custom_kernel(data: input_t) -> output_t:
    """
    rocWMMA最適化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()
    
    # rocWMMA推奨：MFMA最適化パラメータ
    # wave tile fragment sizes（MFMA 16x16x16）
    if M <= 64 and N <= 64:  # 小サイズ：単一wavefront最適化
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 32, 32, 32
        GROUP_SIZE_M = 4
        num_warps = 1  # 単一wavefront
        num_stages = 1
    elif M <= 512 and N <= 512:  # 中サイズ：MFMA 16x16最適化
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 64, 32
        GROUP_SIZE_M = 8
        num_warps = 4  # 4 wavefront per workgroup
        num_stages = 2
    elif K > 4096:  # 大K：メモリバンド幅最適化
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 32, 128, 64
        GROUP_SIZE_M = 8
        num_warps = 4
        num_stages = 3
    else:  # 大サイズ：MFMA throughput最適化
        BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K = 64, 128, 32
        GROUP_SIZE_M = 8
        num_warps = 8
        num_stages = 2
    
    # rocWMMA推奨：最小1024ワークグループ確保
    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
    
    # MFMA occupancy最適化
    if total_workgroups < 512:  # MI300X 304 CU * 2 = 608推奨
        if BLOCK_SIZE_M > 32:
            BLOCK_SIZE_M = 32
        if BLOCK_SIZE_N > 64:
            BLOCK_SIZE_N = 64
    
    # グリッド計算（修正）
    grid = lambda META: (
        triton.cdiv(M, META['BLOCK_SIZE_M']) * triton.cdiv(N, META['BLOCK_SIZE_N']),
    )
    
    # rocWMMA最適化カーネル起動
    fp8_gemm_rocwmma_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_warps=num_warps,
        num_stages=num_stages,
        matrix_instr_nonkdim=16,  # MFMA 16x16（rocWMMA最適）
        waves_per_eu=2,  # VGPR効率最適化
    )
    
    return c

python eval.py benchmark /tmp/tmpra_begr8
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)
