#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpus MI300
from task import input_t, output_t # Assuming input_t is defined elsewhere
import torch
import triton
import triton.language as tl
from typing import Optional, Dict, Tuple, Set
import os
import sys
import subprocess # For running rocprof and ISA dump
import tempfile # For temporary input file
import shutil # For directory cleanup

# --- Set deterministic clock at startup (optional: only do this if on AMD/MI300) ---
# if sys.platform.startswith("linux"):
#     try:
#         subprocess.run(['rocm-smi', '--setperfdeterminism', '2000'], check=False, capture_output=True)
#         print("Set ROCm clock to deterministic 2000MHz.", file=sys.stderr)
#     except Exception as e:
#         print(f"Warning: Could not set ROCm deterministic clock: {e}", file=sys.stderr)

# --- Global state for profiling ---
profile_has_run = False
printed_configs_hardcoded: Set[Tuple[int, int, int]] = set()
# --- Configuration Data (Keep as before) ---
BEST_CONFIGS: Dict[Tuple[int, int, int], Dict] = {
    # (M, N, K): {'BLOCK_SIZE_M': ..., 'BLOCK_SIZE_N': ..., ...}
    (1024, 1536, 7168): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (1024, 3072, 1536): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}, # Note: Log showed GROUP_SIZE_M: 8, but used 4 here as more common. Re-check if score is low.
    (1024, 576,  7168): {'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2},
    (1024, 7168, 256 ): {'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2},
    (1024, 7168, 2048): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 3},
    (1024, 4608, 7168): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (1024, 7168, 2304): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 3},
    (1024, 512,  7168): {'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2},
    (1024, 4096, 512 ): {'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}, # Note: Log showed M=128, G=8. Trying smaller M based on other small N cases. Re-check if low score.
    (6144, 1536, 7168): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 3072, 1536): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 576,  7168): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}, # Note: Log showed M=128, K=64, G=1. Trying smaller M/K based on small N trend. Re-check if low.
    (6144, 7168, 256 ): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 16, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (6144, 7168, 2048): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 4608, 7168): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 7168, 2304): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 512,  7168): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}, # Note: Log showed M=128, K=64, G=1. Trying smaller M/K. Re-check if low.
    (6144, 4096, 512 ): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
}

DEFAULT_CONFIG = {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
# --- Triton Kernel Definition (Identical) ---
@triton.heuristics({ "EVEN_K": lambda args: args["K"] % args["SCALE_BLOCK_K"] == 0, })
@triton.jit
def _gemm_fp8_blockscale_kernel(
    a_ptr, b_ptr, c_ptr, a_scale_ptr, b_scale_ptr, M, N, K,
    stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
    stride_ascale_m, stride_ascale_k, stride_bscale_k, stride_bscale_n,
    SCALE_BLOCK_K: tl.constexpr, SCALE_BLOCK_N: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr,
    EVEN_K: tl.constexpr,
):
    # ... (Kernel code is identical) ...
    # Static asserts
    tl.static_assert(SCALE_BLOCK_K == 128); tl.static_assert(SCALE_BLOCK_N == 128)
    tl.static_assert(SCALE_BLOCK_K % BLOCK_SIZE_K == 0); tl.static_assert(BLOCK_SIZE_N % SCALE_BLOCK_N == 0)
    # PID calculation
    GRID_M = tl.cdiv(M, BLOCK_SIZE_M); GRID_N = tl.cdiv(N, BLOCK_SIZE_N); GRID_MN = GRID_M * GRID_N
    NUM_XCDS: tl.constexpr = 8; pid = tl.program_id(axis=0); num_pid_m = GRID_M; num_pid_n = GRID_N
    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS; tall_xcds = GRID_MN % NUM_XCDS
    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds; xcd = pid % NUM_XCDS; local_pid = pid // NUM_XCDS
    if xcd < tall_xcds: pid = xcd * pids_per_xcd + local_pid
    else: pid = tall_xcds * pids_per_xcd + (xcd - tall_xcds) * (pids_per_xcd - 1) + local_pid
    if GROUP_SIZE_M == 1: pid_m = pid // num_pid_n; pid_n = pid % num_pid_n
    else:
        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
    # Pointers
    offs_m_block = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M); offs_n_block = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k_inner = tl.arange(0, BLOCK_SIZE_K); a_start_ptr = a_ptr + offs_m_block[:, None] * stride_am; b_start_ptr = b_ptr + offs_n_block[None, :] * stride_bn
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    num_k_scale_blocks = tl.cdiv(K, SCALE_BLOCK_K); INNER_K_ITERATIONS = SCALE_BLOCK_K // BLOCK_SIZE_K
    # Loops
    for k_scale_idx in range(num_k_scale_blocks):
        block_accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for k_inner_idx in range(INNER_K_ITERATIONS):
            k_offset = k_scale_idx * SCALE_BLOCK_K + k_inner_idx * BLOCK_SIZE_K; current_k = offs_k_inner + k_offset
            a_ptrs = a_start_ptr + (current_k[None, :] * stride_ak); b_ptrs = b_start_ptr + (current_k[:, None] * stride_bk)
            a_mask = current_k[None, :] < K; b_mask = current_k[:, None] < K
            a = tl.load(a_ptrs, mask=a_mask, other=0.0); b = tl.load(b_ptrs, mask=b_mask, other=0.0)
            block_accumulator += tl.dot(a, b)
        # Scaling
        a_scale_ptr_k = a_scale_ptr + offs_m_block * stride_ascale_m + k_scale_idx * stride_ascale_k; offs_n_scale_block = offs_n_block // SCALE_BLOCK_N
        b_scale_ptr_k = b_scale_ptr + k_scale_idx * stride_bscale_k + offs_n_scale_block * stride_bscale_n; m_mask_scale = offs_m_block < M; n_mask_scale = offs_n_block < N
        a_scale = tl.load(a_scale_ptr_k, mask=m_mask_scale, other=1.0)
        b_scale_load_ptrs = b_scale_ptr + k_scale_idx * stride_bscale_k + (offs_n_block // SCALE_BLOCK_N) * stride_bscale_n
        b_scale = tl.load(b_scale_load_ptrs, mask=n_mask_scale, other=1.0); scaled_block_result = block_accumulator * a_scale[:, None] * b_scale[None, :]
        accumulator += scaled_block_result
    # Write results
    c = accumulator.to(c_ptr.type.element_ty); 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 + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]; c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)
# -------------------------------------------------

# --- Function to Run Kernel (for profiling target - Identical) ---
def _run_kernel_for_profiling(M, N, K):
    # ... (This function remains identical) ...
    print(f"--- Profiling Target Run: M={M}, N={N}, K={K} ---", file=sys.stderr)
    config_key = (M, N, K); selected_config = BEST_CONFIGS.get(config_key, DEFAULT_CONFIG)
    print(f"Using Config: {selected_config}", file=sys.stderr)
    try:
        a = torch.empty((M, K), dtype=torch.float8_e4m3fnuz, device='cuda'); b_in = torch.empty((N, K), dtype=torch.float8_e4m3fnuz, device='cuda'); b = b_in.T
        SCALE_BLOCK_K = 128; SCALE_BLOCK_N = 128
        a_scale = torch.empty((M, triton.cdiv(K, SCALE_BLOCK_K)), dtype=torch.float32, device='cuda')
        b_scale_in = torch.empty((triton.cdiv(N, SCALE_BLOCK_N), triton.cdiv(K, SCALE_BLOCK_K)), dtype=torch.float32, device='cuda'); b_scale = b_scale_in.T
        c = torch.empty((M, N), dtype=torch.bfloat16, device='cuda'); torch.cuda.synchronize()
        print("Dummy data created.", file=sys.stderr)
    except Exception as e: print(f"Error creating dummy data: {e}", file=sys.stderr); return
    grid = (triton.cdiv(M, selected_config['BLOCK_SIZE_M']) * triton.cdiv(N, selected_config['BLOCK_SIZE_N']),)
    try:
        _gemm_fp8_blockscale_kernel[grid](
            a, b, c, a_scale, b_scale, M, N, K,
            a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1),
            a_scale.stride(0), a_scale.stride(1), b_scale.stride(0), b_scale.stride(1),
            SCALE_BLOCK_K=SCALE_BLOCK_K, SCALE_BLOCK_N=SCALE_BLOCK_N,
            BLOCK_SIZE_M=selected_config['BLOCK_SIZE_M'], BLOCK_SIZE_N=selected_config['BLOCK_SIZE_N'],
            BLOCK_SIZE_K=selected_config['BLOCK_SIZE_K'], GROUP_SIZE_M=selected_config['GROUP_SIZE_M'],
            num_warps=selected_config['num_warps'], num_stages=selected_config['num_stages']
        ); torch.cuda.synchronize(); print("Kernel launch and sync complete.", file=sys.stderr)
    except Exception as e: print(f"Error launching kernel for profiling: {e}", file=sys.stderr)

# --- Check if running as profiling target (Identical) ---
_profile_trigger_val = os.environ.get('TRITON_PROFILE_TRIGGER')
if _profile_trigger_val:
    try: _M, _N, _K = map(int, _profile_trigger_val.split(',')); _run_kernel_for_profiling(_M, _N, _K)
    except Exception as e: print(f"Profiling target error: {e}", file=sys.stderr)
    sys.exit(0)
# --------------------------------------------


# --- Main custom_kernel function (Modified Profiling Section) ---
#!POPCORN leaderboard amd-fp8-mm
#!POPCORN gpus MI300
from task import input_t, output_t # Assuming input_t is defined elsewhere
import torch
import triton
import triton.language as tl
from typing import Optional, Dict, Tuple, Set
import os
import sys
import subprocess # For running rocprof and ISA dump
import tempfile # For temporary input file

# --- Set deterministic clock at startup (optional: only do this if on AMD/MI300) ---
# if sys.platform.startswith("linux"):
#     try:
#         subprocess.run(['rocm-smi', '--setperfdeterminism', '2000'], check=False, capture_output=True)
#         print("Set ROCm clock to deterministic 2000MHz.", file=sys.stderr)
#     except Exception as e:
#         print(f"Warning: Could not set ROCm deterministic clock: {e}", file=sys.stderr)

# --- Global state for profiling ---
profile_has_run = False
printed_configs_hardcoded: Set[Tuple[int, int, int]] = set()
# --- Configuration Data (Keep as before) ---
BEST_CONFIGS: Dict[Tuple[int, int, int], Dict] = {
    # (M, N, K): {'BLOCK_SIZE_M': ..., 'BLOCK_SIZE_N': ..., ...}
    (1024, 1536, 7168): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (1024, 3072, 1536): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}, # Note: Log showed GROUP_SIZE_M: 8, but used 4 here as more common. Re-check if score is low.
    (1024, 576,  7168): {'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2},
    (1024, 7168, 256 ): {'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2},
    (1024, 7168, 2048): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 3},
    (1024, 4608, 7168): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (1024, 7168, 2304): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 3},
    (1024, 512,  7168): {'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2},
    (1024, 4096, 512 ): {'BLOCK_SIZE_M': 32,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}, # Note: Log showed M=128, G=8. Trying smaller M based on other small N cases. Re-check if low score.
    (6144, 1536, 7168): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 3072, 1536): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 576,  7168): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}, # Note: Log showed M=128, K=64, G=1. Trying smaller M/K based on small N trend. Re-check if low.
    (6144, 7168, 256 ): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 16, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (6144, 7168, 2048): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 4608, 7168): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 7168, 2304): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (6144, 512,  7168): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}, # Note: Log showed M=128, K=64, G=1. Trying smaller M/K. Re-check if low.
    (6144, 4096, 512 ): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
}

DEFAULT_CONFIG = {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
# --- Triton Kernel Definition (Identical) ---
@triton.heuristics({ "EVEN_K": lambda args: args["K"] % args["SCALE_BLOCK_K"] == 0, })
@triton.jit
def _gemm_fp8_blockscale_kernel(
    a_ptr, b_ptr, c_ptr, a_scale_ptr, b_scale_ptr, M, N, K,
    stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn,
    stride_ascale_m, stride_ascale_k, stride_bscale_k, stride_bscale_n,
    SCALE_BLOCK_K: tl.constexpr, SCALE_BLOCK_N: tl.constexpr,
    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr,
    EVEN_K: tl.constexpr,
):
    # ... (Kernel code is identical) ...
    # Static asserts
    tl.static_assert(SCALE_BLOCK_K == 128); tl.static_assert(SCALE_BLOCK_N == 128)
    tl.static_assert(SCALE_BLOCK_K % BLOCK_SIZE_K == 0); tl.static_assert(BLOCK_SIZE_N % SCALE_BLOCK_N == 0)
    # PID calculation
    GRID_M = tl.cdiv(M, BLOCK_SIZE_M); GRID_N = tl.cdiv(N, BLOCK_SIZE_N); GRID_MN = GRID_M * GRID_N
    NUM_XCDS: tl.constexpr = 8; pid = tl.program_id(axis=0); num_pid_m = GRID_M; num_pid_n = GRID_N
    pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS; tall_xcds = GRID_MN % NUM_XCDS
    tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds; xcd = pid % NUM_XCDS; local_pid = pid // NUM_XCDS
    if xcd < tall_xcds: pid = xcd * pids_per_xcd + local_pid
    else: pid = tall_xcds * pids_per_xcd + (xcd - tall_xcds) * (pids_per_xcd - 1) + local_pid
    if GROUP_SIZE_M == 1: pid_m = pid // num_pid_n; pid_n = pid % num_pid_n
    else:
        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
    # Pointers
    offs_m_block = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M); offs_n_block = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k_inner = tl.arange(0, BLOCK_SIZE_K); a_start_ptr = a_ptr + offs_m_block[:, None] * stride_am; b_start_ptr = b_ptr + offs_n_block[None, :] * stride_bn
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    num_k_scale_blocks = tl.cdiv(K, SCALE_BLOCK_K); INNER_K_ITERATIONS = SCALE_BLOCK_K // BLOCK_SIZE_K
    # Loops
    for k_scale_idx in range(num_k_scale_blocks):
        block_accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
        for k_inner_idx in range(INNER_K_ITERATIONS):
            k_offset = k_scale_idx * SCALE_BLOCK_K + k_inner_idx * BLOCK_SIZE_K; current_k = offs_k_inner + k_offset
            a_ptrs = a_start_ptr + (current_k[None, :] * stride_ak); b_ptrs = b_start_ptr + (current_k[:, None] * stride_bk)
            a_mask = current_k[None, :] < K; b_mask = current_k[:, None] < K
            a = tl.load(a_ptrs, mask=a_mask, other=0.0); b = tl.load(b_ptrs, mask=b_mask, other=0.0)
            block_accumulator += tl.dot(a, b)
        # Scaling
        a_scale_ptr_k = a_scale_ptr + offs_m_block * stride_ascale_m + k_scale_idx * stride_ascale_k; offs_n_scale_block = offs_n_block // SCALE_BLOCK_N
        b_scale_ptr_k = b_scale_ptr + k_scale_idx * stride_bscale_k + offs_n_scale_block * stride_bscale_n; m_mask_scale = offs_m_block < M; n_mask_scale = offs_n_block < N
        a_scale = tl.load(a_scale_ptr_k, mask=m_mask_scale, other=1.0)
        b_scale_load_ptrs = b_scale_ptr + k_scale_idx * stride_bscale_k + (offs_n_block // SCALE_BLOCK_N) * stride_bscale_n
        b_scale = tl.load(b_scale_load_ptrs, mask=n_mask_scale, other=1.0); scaled_block_result = block_accumulator * a_scale[:, None] * b_scale[None, :]
        accumulator += scaled_block_result
    # Write results
    c = accumulator.to(c_ptr.type.element_ty); 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 + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]; c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
    tl.store(c_ptrs, c, mask=c_mask)
# -------------------------------------------------

# --- Function to Run Kernel (for profiling target - Identical) ---
def _run_kernel_for_profiling(M, N, K):
    # ... (This function remains identical) ...
    print(f"--- Profiling Target Run: M={M}, N={N}, K={K} ---", file=sys.stderr)
    config_key = (M, N, K); selected_config = BEST_CONFIGS.get(config_key, DEFAULT_CONFIG)
    print(f"Using Config: {selected_config}", file=sys.stderr)
    try:
        a = torch.empty((M, K), dtype=torch.float8_e4m3fnuz, device='cuda'); b_in = torch.empty((N, K), dtype=torch.float8_e4m3fnuz, device='cuda'); b = b_in.T
        SCALE_BLOCK_K = 128; SCALE_BLOCK_N = 128
        a_scale = torch.empty((M, triton.cdiv(K, SCALE_BLOCK_K)), dtype=torch.float32, device='cuda')
        b_scale_in = torch.empty((triton.cdiv(N, SCALE_BLOCK_N), triton.cdiv(K, SCALE_BLOCK_K)), dtype=torch.float32, device='cuda'); b_scale = b_scale_in.T
        c = torch.empty((M, N), dtype=torch.bfloat16, device='cuda'); torch.cuda.synchronize()
        print("Dummy data created.", file=sys.stderr)
    except Exception as e: print(f"Error creating dummy data: {e}", file=sys.stderr); return
    grid = (triton.cdiv(M, selected_config['BLOCK_SIZE_M']) * triton.cdiv(N, selected_config['BLOCK_SIZE_N']),)
    try:
        _gemm_fp8_blockscale_kernel[grid](
            a, b, c, a_scale, b_scale, M, N, K,
            a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1),
            a_scale.stride(0), a_scale.stride(1), b_scale.stride(0), b_scale.stride(1),
            SCALE_BLOCK_K=SCALE_BLOCK_K, SCALE_BLOCK_N=SCALE_BLOCK_N,
            BLOCK_SIZE_M=selected_config['BLOCK_SIZE_M'], BLOCK_SIZE_N=selected_config['BLOCK_SIZE_N'],
            BLOCK_SIZE_K=selected_config['BLOCK_SIZE_K'], GROUP_SIZE_M=selected_config['GROUP_SIZE_M'],
            num_warps=selected_config['num_warps'], num_stages=selected_config['num_stages']
        ); torch.cuda.synchronize(); print("Kernel launch and sync complete.", file=sys.stderr)
    except Exception as e: print(f"Error launching kernel for profiling: {e}", file=sys.stderr)

# --- Check if running as profiling target (Identical) ---
_profile_trigger_val = os.environ.get('TRITON_PROFILE_TRIGGER')
if _profile_trigger_val:
    try: _M, _N, _K = map(int, _profile_trigger_val.split(',')); _run_kernel_for_profiling(_M, _N, _K)
    except Exception as e: print(f"Profiling target error: {e}", file=sys.stderr)
    sys.exit(0)
# --------------------------------------------


# --- Main custom_kernel function (Modified Profiling Section) ---
def custom_kernel(data: input_t) -> output_t:
    """
    Computes FP8 GEMM, potentially triggers rocprofv2 once with specific cache metrics using an input file.
    """
    global profile_has_run, printed_configs_hardcoded
    a, b_in, a_scale, b_scale_in, c = data
    # Assertions & Transpose (Identical)
    assert a.device.type == 'cuda'; assert b_in.device.type == 'cuda'; assert a_scale.device.type == 'cuda'; assert b_scale_in.device.type == 'cuda'; assert c.device.type == 'cuda'; assert a.dtype == torch.float8_e4m3fnuz; assert b_in.dtype == torch.float8_e4m3fnuz; assert a_scale.dtype == torch.float32; assert b_scale_in.dtype == torch.float32; assert c.dtype == torch.bfloat16
    b = b_in.T; b_scale = b_scale_in.T
    M, K = a.shape; K_b, N = b.shape; M_c, N_c = c.shape
    assert K == K_b; assert M == M_c; assert N == N_c
    SCALE_BLOCK_K = 128; SCALE_BLOCK_N = 128; expected_a_scale_shape = (M, triton.cdiv(K, SCALE_BLOCK_K)); expected_b_scale_shape = (triton.cdiv(K, SCALE_BLOCK_K), triton.cdiv(N, SCALE_BLOCK_N)); assert a_scale.shape == expected_a_scale_shape; assert b_scale.shape == expected_b_scale_shape
    # Select Config (Identical)
    config_key = (M, N, K); selected_config = BEST_CONFIGS.get(config_key, DEFAULT_CONFIG); using_default = config_key not in BEST_CONFIGS

    # --- Trigger Profiling Once ---
    if not profile_has_run:
        profile_has_run = True
        print(f"--- Attempting rocprofv2 run for M={M}, N={N}, K={K} with CACHE METRICS via input file ---", file=sys.stderr)
        input_counters_file = "rocprofv2_counters.txt"
        files_to_cleanup = [input_counters_file]

        try:
            # --- First run rocprofv2 --list-counters (Optional but good for debug) ---
            print("--- Running rocprofv2 --list-counters ---", file=sys.stderr)
            try:
                counters_cmd = ['rocprofv2', '--list-counters']
                counters_result = subprocess.run(counters_cmd, capture_output=True, text=True, check=False)
                print(f"--- rocprofv2 --list-counters exit code: {counters_result.returncode} ---", file=sys.stderr)
            except FileNotFoundError:
                print("ERROR: rocprofv2 command not found.", file=sys.stderr)
            except Exception as e:
                print(f"ERROR running rocprofv2 --list-counters: {e}", file=sys.stderr)

            # --- Create the input counters file ---
            counters_to_profile = [
                "VALUBusy",
                "MfmaUtil",
                "WaveDepWait",
                "WaveIssueWait",
                "LDSBankConflict",
                "MemUnitBusy",
                "L2CacheHitRate",
                "ValuPipeIssueUtil"
            ]
            try:
                with open(input_counters_file, 'w') as f:
                    f.write("pmc: " + ' '.join(counters_to_profile))
                print(f"Created counter input file: {input_counters_file} with reduced GRBM counters", file=sys.stderr)
            except Exception as e:
                print(f"ERROR creating counter input file {input_counters_file}: {e}", file=sys.stderr)
                raise

            # --- Set env var & build command ---
            os.environ['TRITON_PROFILE_TRIGGER'] = f"{M},{N},{K}"; script_path = os.path.abspath(__file__)
            os.environ['HIP_FORCE_DEV_KERNARG'] = '1'
            output_dir = "."
            default_output_csv = os.path.join(output_dir, "results.csv")
            cmd_rocprof = ['rocprofv2', '-i', input_counters_file, '-d', output_dir, sys.executable, script_path]
            print(f"Running command: {' '.join(cmd_rocprof)}", file=sys.stderr)

            # --- Execute ---
            timeout_seconds = 120
            result_rocprof = subprocess.run(cmd_rocprof, capture_output=True, text=True, check=False, timeout=timeout_seconds)

            # --- Print diagnostics ---
            print("--- rocprofv2 execution stdout ---", file=sys.stderr); print(result_rocprof.stdout, file=sys.stderr)
            print("--- rocprofv2 execution stderr ---", file=sys.stderr); print(result_rocprof.stderr, file=sys.stderr)
            print(f"--- rocprofv2 exit code: {result_rocprof.returncode} ---", file=sys.stderr)

            # --- Read and print CSV (looking for default name) ---
            if result_rocprof.returncode == 0:
                print(f"--- Contents of {default_output_csv} ---", file=sys.stderr)
                try:
                    with open(default_output_csv, 'r') as f: print(f.read(), file=sys.stderr)
                except FileNotFoundError: print(f"Warning: Output file {default_output_csv} not found.", file=sys.stderr)
                except Exception as e: print(f"Error reading output file {default_output_csv}: {e}", file=sys.stderr)
            else: print("rocprofv2 failed or exited non-zero, not reading output file.", file=sys.stderr)

        except FileNotFoundError: print("ERROR: rocprofv2 or python command not found.", file=sys.stderr)
        except subprocess.TimeoutExpired: print(f"ERROR: rocprofv2 timed out after {timeout_seconds} seconds.", file=sys.stderr)
        except Exception as e: print(f"ERROR running rocprofv2: {e}", file=sys.stderr)
        finally:
            # --- Cleanup Env Vars and Restore Clock ---
            if 'TRITON_PROFILE_TRIGGER' in os.environ: del os.environ['TRITON_PROFILE_TRIGGER']
            if 'HIP_FORCE_DEV_KERNARG' in os.environ: del os.environ['HIP_FORCE_DEV_KERNARG']
            if sys.platform.startswith("linux"):
                try:
                    subprocess.run(['rocm-smi', '-r'], check=False, capture_output=True)
                    print("Restored ROCm clock and settings to default with 'rocm-smi -r'.", file=sys.stderr)
                except Exception as e:
                    print(f"Warning: Could not restore ROCm default settings: {e}", file=sys.stderr)

            # --- General Cleanup ---
            print("--- Cleaning up temporary profiling files ---", file=sys.stderr)
            if os.path.exists(default_output_csv):
                files_to_cleanup.append(default_output_csv)
            else:
                print(f"Note: Default output file {default_output_csv} not found for cleanup.", file=sys.stderr)

            for fname in files_to_cleanup:
                try: os.remove(fname); print(f"Removed {fname}", file=sys.stderr)
                except OSError as e: print(f"Warning: Could not remove {fname}: {e}", file=sys.stderr)
            print("--- End profiling attempt ---", file=sys.stderr)
    # --- End Profiling Trigger ---

    # --- Grid Calculation (Identical) ---
    grid = (triton.cdiv(M, selected_config['BLOCK_SIZE_M']) * triton.cdiv(N, selected_config['BLOCK_SIZE_N']),)
    # --- Launch Kernel (Identical) ---
    _gemm_fp8_blockscale_kernel[grid](
        a, b, c, a_scale, b_scale, M, N, K,
        a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1),
        a_scale.stride(0), a_scale.stride(1), b_scale.stride(0), b_scale.stride(1),
        SCALE_BLOCK_K=SCALE_BLOCK_K, SCALE_BLOCK_N=SCALE_BLOCK_N,
        BLOCK_SIZE_M=selected_config['BLOCK_SIZE_M'], BLOCK_SIZE_N=selected_config['BLOCK_SIZE_N'],
        BLOCK_SIZE_K=selected_config['BLOCK_SIZE_K'], GROUP_SIZE_M=selected_config['GROUP_SIZE_M'],
        num_warps=selected_config['num_warps'], num_stages=selected_config['num_stages']
    )
    return c

# --- (Keep the standalone __main__ block if desired) ---
# if __name__ == '__main__':
#    pass



python eval.py benchmark /tmp/tmprv4i52a_
--- Attempting rocprofv2 run for M=1024, N=1536, K=7168 with CACHE METRICS via input file ---
--- Running rocprofv2 --list-counters ---
--- rocprofv2 --list-counters exit code: 1 ---
Created counter input file: rocprofv2_counters.txt with reduced GRBM counters
Running command: rocprofv2 -i rocprofv2_counters.txt -d . /home/runner/_work/_tool/Python/3.10.17/x64/bin/python /home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py
--- rocprofv2 execution stdout ---
ROCProfilerV2: Collecting the following counters:
- VALUBusy
- MfmaUtil
- WaveDepWait
- WaveIssueWait
- LDSBankConflict
- MemUnitBusy
- L2CacheHitRate
- ValuPipeIssueUtil
Enabling Counter Collection
Results File: "./pmc_1/results_655.csv"

The output path for the following counters: ./pmc_1

--- rocprofv2 execution stderr ---
--- Profiling Target Run: M=1024, N=1536, K=7168 ---
Using Config: {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
Dummy data created.
Kernel launch and sync complete.

--- rocprofv2 exit code: 0 ---
--- Contents of ./results.csv ---
Warning: Output file ./results.csv not found.
Restored ROCm clock and settings to default with 'rocm-smi -r'.
--- Cleaning up temporary profiling files ---
Note: Default output file ./results.csv not found for cleanup.
Removed rocprofv2_counters.txt
--- End profiling attempt ---
