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 benchmark /tmp/tmp2v4a8jfk
/usr/local/lib/python3.10/dist-packages/torch/_inductor/compile_fx.py:246: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn(
E0521 09:05:53.292000 8717 torch/_inductor/select_algorithm.py:2148] [0/0] Exception No valid triton configs. OutOfResources: out of resource: shared memory, Required: 131072, Hardware limit: 65536. Reducing block sizes or `num_stages` may help. for benchmark choice TritonTemplateCaller(/tmp/torchinductor_runner/l7/cl7v27266igh5eoqkmuychai47wcmpll4ekmc6jowtlvgzeqgvjo.py, ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8)
E0521 09:05:54.511000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Runtime error during autotuning: 
E0521 09:05:54.511000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 131072, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:05:54.511000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Ignoring this choice.
E0521 09:05:54.925000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Runtime error during autotuning: 
E0521 09:05:54.925000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:05:54.925000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Ignoring this choice.
E0521 09:05:55.588000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Runtime error during autotuning: 
E0521 09:05:55.588000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:05:55.588000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Ignoring this choice.
E0521 09:05:56.389000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Runtime error during autotuning: 
E0521 09:05:56.389000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 131072, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:05:56.389000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Ignoring this choice.
E0521 09:05:56.648000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Runtime error during autotuning: 
E0521 09:05:56.648000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:05:56.648000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Ignoring this choice.
E0521 09:05:56.768000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Runtime error during autotuning: 
E0521 09:05:56.768000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 81920, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:05:56.768000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Ignoring this choice.
E0521 09:05:57.167000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Runtime error during autotuning: 
E0521 09:05:57.167000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:05:57.167000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Ignoring this choice.
E0521 09:05:57.284000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Runtime error during autotuning: 
E0521 09:05:57.284000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 131072, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:05:57.284000 8717 torch/_inductor/select_algorithm.py:2274] [0/0] Ignoring this choice.
AUTOTUNE mm(1024x7168, 7168x1536)
strides: [7168, 1], [1, 7168]
dtypes: torch.float32, torch.float32
  triton_mm_6 0.3114 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_15 0.3178 ms 98.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_16 0.3205 ms 97.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_18 0.3243 ms 96.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_10 0.3296 ms 94.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=16, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_11 0.3352 ms 92.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_8 0.3375 ms 92.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_7 0.3412 ms 91.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_4 0.3450 ms 90.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=32, BLOCK_N=64, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_14 0.3583 ms 86.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
SingleProcess AUTOTUNE benchmarking takes 3.6229 seconds and 0.7116 seconds precompiling for 37 choices
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) )

E0521 09:06:07.727000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:06:07.727000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 131072, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:06:07.727000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:06:08.219000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:06:08.219000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:06:08.219000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:06:08.838000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:06:08.838000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:06:08.838000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:06:09.577000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:06:09.577000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 131072, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:06:09.577000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:06:09.831000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:06:09.831000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:06:09.831000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:06:09.948000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:06:09.948000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 81920, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:06:09.948000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:06:10.335000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:06:10.335000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:06:10.335000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:06:10.453000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:06:10.453000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] No valid triton configs. OutOfResources: out of resource: shared memory, Required: 131072, Hardware limit: 65536. Reducing block sizes or `num_stages` may help.. 
E0521 09:06:10.453000 8717 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
AUTOTUNE mm(1024x1536, 1536x3072)
strides: [128*s23, 1], [1, 128*s23]
dtypes: torch.float32, torch.float32
  mm 0.1238 ms 100.0% 
  triton_mm_46 0.1321 ms 93.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=16, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_63 0.1353 ms 91.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_50 0.1354 ms 91.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=128, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_54 0.1374 ms 90.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_47 0.1399 ms 88.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_56 0.1424 ms 86.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=64, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=2, num_stages=2, num_warps=4
  triton_mm_42 0.1426 ms 86.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_43 0.1445 ms 85.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=32, EVEN_K=True, GROUP_M=8, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
  triton_mm_61 0.1473 ms 84.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=128, BLOCK_N=128, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=1, num_stages=2, num_warps=8
SingleProcess AUTOTUNE benchmarking takes 4.4571 seconds and 1.0326 seconds precompiling for 37 choices
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) )

[...] 13 lines omitted