import os
import torch
from task import input_t, output_t

# 環境変数を設定
# ただし、これらの設定は、PyTorchやROCmが初期化される前に有効である必要があるため、
# スクリプトの実行方法によっては効果がない場合があります。
# シェルで設定する方が確実です。
os.environ['TORCH_NCCL_USE_TENSOR_REGISTER_ALLOCATOR_HOOK'] = '1'
os.environ['HIP_FORCE_DEV_KERNARG'] = '1'
os.environ['RCCL_MSCCLPP_THRESHOLD'] = str(2*1024*1024*1024) # Pythonで計算して文字列で設定
os.environ['MSCCLPP_READ_ALLRED'] = '1'
os.environ['ROCR_VISIBLE_DEVICES'] = '0,1,2,3,4,5,6,7'

# @torch.compile の最適化デコレータ
@torch.compile(mode="max-autotune")
def custom_kernel(data: input_t) -> output_t:
    """
    MI300X向けに最適化されたブロックスケールfp8 GEMMの実装。
    PyTorchのtorch.compileを使用して既存のロジックを最適化します。
    """
    a, b, a_scale, b_scale, c = data

    a = a.contiguous()
    a_scale = a_scale.contiguous()
    b_scale = b_scale.contiguous()

    m = a.shape[0]
    n = b.shape[0]
    k = a.shape[1]
    block_shape_n = 128
    block_shape_k = 128
    scale_n = b_scale.shape[0]
    scale_k = b_scale.shape[1]

    a_scale = a_scale.unsqueeze(-1).repeat(1, 1, block_shape_k)
    a_scale = a_scale.reshape(m, scale_k * block_shape_k)
    a_scale = a_scale[:, :k]

    a = a.to(a_scale.dtype) * a_scale

    b_scale = (
        b_scale.view(-1, 1)
        .repeat(1, block_shape_n * block_shape_k)
        .view(scale_n, scale_k, block_shape_n, block_shape_k)
        .permute(0, 2, 1, 3)
        .reshape(scale_n * block_shape_n, scale_k * block_shape_k)
    )
    b_scale = b_scale[:n, :k]

    b = b.to(b_scale.dtype) * b_scale

    c.copy_( (a @ b.T).to(torch.bfloat16) )
    return c
python eval.py leaderboard /tmp/tmpgig7d0h_
skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

skipping cudagraphs due to mutated inputs (1 instances). Found from : 
   File "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py", line 53, in custom_kernel
    c.copy_( (a @ b.T).to(torch.bfloat16) )

