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 test /tmp/tmp0newa5vl
/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(
AUTOTUNE mm(64x128, 128x64)
strides: [128, 1], [1, 128]
dtypes: torch.float32, torch.float32
  triton_mm_22 0.0068 ms 100.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, 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=8
  triton_mm_20 0.0069 ms 98.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_15 0.0070 ms 97.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, 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=8
  triton_mm_17 0.0077 ms 88.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, 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=2, num_stages=2, num_warps=4
  triton_mm_16 0.0081 ms 83.7% 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=8
  triton_mm_12 0.0083 ms 82.5% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_14 0.0096 ms 71.1% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, 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=2, num_stages=2, num_warps=4
  triton_mm_21 0.0101 ms 67.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=32, BLOCK_M=64, BLOCK_N=64, 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
  triton_mm_2 0.0101 ms 67.7% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=16, BLOCK_M=32, BLOCK_N=32, 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_18 0.0103 ms 66.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=4
SingleProcess AUTOTUNE benchmarking takes 1.9003 seconds and 0.2278 seconds precompiling for 24 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:04:39.875000 562 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:04:39.875000 562 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:04:39.875000 562 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:04:40.392000 562 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:04:40.392000 562 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:04:40.392000 562 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:04:41.556000 562 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:04:41.556000 562 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:04:41.556000 562 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
E0521 09:04:41.809000 562 torch/_inductor/select_algorithm.py:2274] [0/1] Runtime error during autotuning: 
E0521 09:04:41.809000 562 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:04:41.809000 562 torch/_inductor/select_algorithm.py:2274] [0/1] Ignoring this choice.
AUTOTUNE mm(64x7168, 7168x1536)
strides: [128*s23, 1], [1, 128*s23]
dtypes: torch.float32, torch.float32
  mm 0.0408 ms 100.0% 
  triton_mm_24 0.0887 ms 46.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=256, BLOCK_M=32, BLOCK_N=16, EVEN_K=False, GROUP_M=4, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=2
  triton_mm_26 0.0940 ms 43.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=32, 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_31 0.0955 ms 42.8% 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_28 0.1074 ms 38.0% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=16, 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_30 0.1252 ms 32.6% 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_27 0.1265 ms 32.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_25 0.1415 ms 28.8% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=16, BLOCK_M=32, BLOCK_N=32, 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_32 0.1537 ms 26.6% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, 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_29 0.1538 ms 26.5% 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
SingleProcess AUTOTUNE benchmarking takes 3.8176 seconds and 0.9464 seconds precompiling for 31 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) )

E0521 09:04:59.076000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Runtime error during autotuning: 
E0521 09:04:59.076000 562 torch/_inductor/select_algorithm.py:2274] [0/2] 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:04:59.076000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Ignoring this choice.
E0521 09:04:59.557000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Runtime error during autotuning: 
E0521 09:04:59.557000 562 torch/_inductor/select_algorithm.py:2274] [0/2] 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:04:59.557000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Ignoring this choice.
E0521 09:05:00.156000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Runtime error during autotuning: 
E0521 09:05:00.156000 562 torch/_inductor/select_algorithm.py:2274] [0/2] 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:00.156000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Ignoring this choice.
E0521 09:05:00.876000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Runtime error during autotuning: 
E0521 09:05:00.876000 562 torch/_inductor/select_algorithm.py:2274] [0/2] 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:00.876000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Ignoring this choice.
E0521 09:05:01.116000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Runtime error during autotuning: 
E0521 09:05:01.116000 562 torch/_inductor/select_algorithm.py:2274] [0/2] 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:01.116000 562 torch/_inductor/select_algorithm.py:2274] [0/2] Ignoring this choice.
AUTOTUNE mm(96x256, 256x7168)
strides: [128*s23, 1], [1, 128*s23]
dtypes: torch.float32, torch.float32
  mm 0.0135 ms 100.0% 
  triton_mm_59 0.0139 ms 97.4% 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_55 0.0140 ms 96.3% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=16, BLOCK_M=32, BLOCK_N=32, 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_61 0.0145 ms 93.1% 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_60 0.0146 ms 92.6% 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_71 0.0148 ms 91.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_65 0.0150 ms 90.4% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=128, BLOCK_M=64, BLOCK_N=64, EVEN_K=True, GROUP_M=16, USE_FAST_ACCUM=False, kpack=2, matrix_instr_nonkdim=16, waves_per_eu=0, num_stages=2, num_warps=8
  triton_mm_72 0.0150 ms 89.9% ACC_TYPE='tl.float32', ALLOW_TF32=False, BLOCK_K=64, 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_63 0.0151 ms 89.6% 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_57 0.0152 ms 88.7% 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
SingleProcess AUTOTUNE benchmarking takes 4.1144 seconds and 0.9669 seconds precompiling for 35 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) )

