#!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) ---
ALT_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': 8, 'num_warps': 8, 'num_stages': 2},
    (1024, 576,  7168): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (1024, 7168, 256 ): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (1024, 7168, 2048): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, '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': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (1024, 512,  7168): {'BLOCK_SIZE_M': 64,  'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (1024, 4096, 512 ): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 8, 'num_warps': 8, 'num_stages': 2},
    (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': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 1, 'num_warps': 8, 'num_stages': 2},
    (6144, 7168, 256 ): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 3},
    (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': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 1, 'num_warps': 8, 'num_stages': 2},
    (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},
}

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': 2},
    (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': 2},
    (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': 2},
    (6144, 3072, 1536): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (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': 2},
    (6144, 4608, 7168): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (6144, 7168, 2304): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2},
    (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': 2},
}

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 rocprof once with specific metrics
    by using multiple rocprof runs with metric groups.
    """
    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 rocprof runs for M={M}, N={N}, K={K} with grouped metrics ---", file=sys.stderr)
        
        # Define metric groups
        metric_groups = [
            {
                "name": "occupancy_stalls",
                "metrics": "MeanOccupancyPerCU MeanOccupancyPerActiveCU ALUStalledByLDS MemUnitBusy MemUnitStalled Wavefronts",
                "output": "rocprof_metrics_group1.csv",
                "input": "rocprof_input_group1.txt"
            },
            {
                "name": "compute_throughput",
                "metrics": "VALUBusy VALUUtilization SQ_INSTS_MFMA FetchSize WriteSize SQ_WAVES",
                "output": "rocprof_metrics_group2.csv",
                "input": "rocprof_input_group2.txt"
            }
        ]
        
        script_path = os.path.abspath(__file__)
        files_to_cleanup = []
        
        # Log counter availability
        # try:
        #     print("--- Running rocprofv2 --list-counters ---", file=sys.stderr)
        #     counters_cmd = ['rocprofv2', '--list-counters']
        #     counters_result = subprocess.run(counters_cmd, capture_output=True, text=True, check=False)
        #     print("--- rocprofv2 --list-counters stdout ---", file=sys.stderr)
        #     print(counters_result.stdout, file=sys.stderr)
        #     print("--- rocprofv2 --list-counters stderr ---", file=sys.stderr)
        #     print(counters_result.stderr, file=sys.stderr)
        #     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)
            
        # Run rocprof for each metric group
        for group in metric_groups:
            print(f"--- Running rocprof for {group['name']} metrics ---", file=sys.stderr)
            input_txt = group['input']
            output_csv = group['output']
            files_to_cleanup.extend([input_txt, output_csv])
            
            try:
                # Create input file for this group
                print(f"Creating {input_txt} with: pmc : {group['metrics']}", file=sys.stderr)
                with open(input_txt, 'w') as f:
                    f.write(f"pmc : {group['metrics']}")
                
                # Set environment variables
                os.environ['TRITON_PROFILE_TRIGGER'] = f"{M},{N},{K}"
                os.environ['HIP_FORCE_DEV_KERNARG'] = '1'
                
                # Build and run rocprof command
                cmd_rocprof = ['rocprof', '--timestamp', 'on', '-i', input_txt, '-o', output_csv, 
                              sys.executable, script_path]
                print(f"Running command: {' '.join(cmd_rocprof)}", file=sys.stderr)
                
                # Execute with timeout
                timeout_seconds = 90
                result_rocprof = subprocess.run(cmd_rocprof, capture_output=True, text=True, 
                                              check=False, timeout=timeout_seconds)
                
                # Print diagnostics
                print(f"--- rocprof {group['name']} execution stdout ---", file=sys.stderr)
                print(result_rocprof.stdout, file=sys.stderr)
                print(f"--- rocprof {group['name']} execution stderr ---", file=sys.stderr)
                print(result_rocprof.stderr, file=sys.stderr)
                print(f"--- rocprof {group['name']} exit code: {result_rocprof.returncode} ---", file=sys.stderr)
                
                # Read and print CSV
                if result_rocprof.returncode == 0:
                    print(f"--- Contents of {output_csv} ---", file=sys.stderr)
                    try:
                        with open(output_csv, 'r') as f:
                            print(f.read(), file=sys.stderr)
                    except FileNotFoundError:
                        print(f"Warning: Output file {output_csv} not found.", file=sys.stderr)
                    except Exception as e:
                        print(f"Error reading output file {output_csv}: {e}", file=sys.stderr)
                else:
                    print(f"rocprof {group['name']} run failed, not reading output file.", file=sys.stderr)
                    
            except FileNotFoundError:
                print("ERROR: rocprof or python command not found.", file=sys.stderr)
            except subprocess.TimeoutExpired:
                print(f"ERROR: rocprof {group['name']} timed out after {timeout_seconds} seconds.", file=sys.stderr)
            except Exception as e:
                print(f"ERROR running rocprof {group['name']}: {e}", file=sys.stderr)
            finally:
                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)

        # --- Part 2: Run for ISA Dump (Same as before) ---
        # print(f"--- Attempting ISA Dump for M={M}, N={N}, K={K} ---", file=sys.stderr)
        # isa_dump_env = os.environ.copy(); isa_dump_env['AMDGCN_ENABLE_DUMP'] = '1'; isa_dump_env['TRITON_PROFILE_TRIGGER'] = f"{M},{N},{K}"
        # isa_dump_env['HIP_FORCE_DEV_KERNARG'] = '1' 
        # try:
        #     cmd_isa = [ sys.executable, script_path ]; print(f"Running command: {' '.join(cmd_isa)} with AMDGCN_ENABLE_DUMP=1", file=sys.stderr)
        #     timeout_seconds_isa = 60; result_isa = subprocess.run( cmd_isa, env=isa_dump_env, capture_output=True, text=True, check=False, timeout=timeout_seconds_isa )
        #     print("--- ISA Dump stdout ---", file=sys.stderr); print(result_isa.stdout, file=sys.stderr)
        #     print("--- ISA Dump stderr ---", file=sys.stderr); print(result_isa.stderr, file=sys.stderr); print(f"--- ISA Dump exit code: {result_isa.returncode} ---", file=sys.stderr)
        # except FileNotFoundError: print("ERROR: python command not found for ISA dump.", file=sys.stderr)
        # except subprocess.TimeoutExpired: print(f"ERROR: ISA dump timed out after {timeout_seconds_isa} seconds.", file=sys.stderr)
        # except Exception as e: print(f"ERROR running ISA dump: {e}", file=sys.stderr)
        # finally:
        #      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)
        for fname in files_to_cleanup:
            try: os.remove(fname)
            except OSError: pass
        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']
    )
    # --- Log selected config (Identical) ---
    if config_key not in printed_configs_hardcoded:
        status = "DEFAULT" if using_default else "SPECIFIC"; print(f"HARDCODED_CONFIG Used ({status} for M={M}, N={N}, K={K}): {selected_config}", file=sys.stderr); printed_configs_hardcoded.add(config_key)

    return c

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

python eval.py benchmark /tmp/tmpgqjw8tzk
--- Attempting rocprof runs for M=1024, N=1536, K=7168 with grouped metrics ---
--- Running rocprof for occupancy_stalls metrics ---
Creating rocprof_input_group1.txt with: pmc : MeanOccupancyPerCU MeanOccupancyPerActiveCU ALUStalledByLDS MemUnitBusy MemUnitStalled Wavefronts
Running command: rocprof --timestamp on -i rocprof_input_group1.txt -o rocprof_metrics_group1.csv /home/runner/_work/_tool/Python/3.10.17/x64/bin/python /home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py
--- rocprof occupancy_stalls execution stdout ---
RPL: on '250423_052529' from '/opt/rocm-6.3.1' in '/home/runner/_work/discord-cluster-manager/discord-cluster-manager'
RPL: profiling '"/home/runner/_work/_tool/Python/3.10.17/x64/bin/python" "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py"'
RPL: input file 'rocprof_input_group1.txt'
RPL: output dir '/tmp/rpl_data_250423_052529_723'
RPL: result dir '/tmp/rpl_data_250423_052529_723/input0_results_250423_052529'
--- 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.
ROCProfiler: input from "/tmp/rpl_data_250423_052529_723/input0.xml"
  gpu_index = 
  kernel = 
  range = 
  6 metrics
    MeanOccupancyPerCU, MeanOccupancyPerActiveCU, ALUStalledByLDS, MemUnitBusy, MemUnitStalled, Wavefronts

ROCPRofiler: 1 contexts collected, output directory /tmp/rpl_data_250423_052529_723/input0_results_250423_052529
File 'rocprof_metrics_group1.csv' is generating

--- rocprof occupancy_stalls execution stderr ---

--- rocprof occupancy_stalls exit code: 0 ---
--- Contents of rocprof_metrics_group1.csv ---
Index,KernelName,gpu-id,queue-id,queue-index,pid,tid,grd,wgr,lds,scr,arch_vgpr,accum_vgpr,sgpr,wave_size,sig,obj,MeanOccupancyPerCU,MeanOccupancyPerActiveCU,ALUStalledByLDS,MemUnitBusy,MemUnitStalled,Wavefronts,DispatchNs,BeginNs,EndNs,CompleteNs
0,"_gemm_fp8_blockscale_kernel.kd",2,0,0,787,787,98304,512,12288,0,112,0,48,64,0x0,0x7f721daf0740,4.7328222604,5.0463335444,0.2445007744,5.1985400621,0.0425262799,1536.0000000000,994553713432127,994553730579684,994553730706371,994553730716572

Restored ROCm clock and settings to default with 'rocm-smi -r'.
--- Running rocprof for compute_throughput metrics ---
Creating rocprof_input_group2.txt with: pmc : VALUBusy VALUUtilization SQ_INSTS_MFMA FetchSize WriteSize SQ_WAVES
Running command: rocprof --timestamp on -i rocprof_input_group2.txt -o rocprof_metrics_group2.csv /home/runner/_work/_tool/Python/3.10.17/x64/bin/python /home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py
--- rocprof compute_throughput execution stdout ---
RPL: on '250423_052539' from '/opt/rocm-6.3.1' in '/home/runner/_work/discord-cluster-manager/discord-cluster-manager'
RPL: profiling '"/home/runner/_work/_tool/Python/3.10.17/x64/bin/python" "/home/runner/_work/discord-cluster-manager/discord-cluster-manager/submission.py"'
RPL: input file 'rocprof_input_group2.txt'
RPL: output dir '/tmp/rpl_data_250423_052539_921'
RPL: result dir '/tmp/rpl_data_250423_052539_921/input0_results_250423_052539'
--- 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.
ROCProfiler: input from "/tmp/rpl_data_250423_052539_921/input0.xml"
  gpu_index = 
  kernel = 
  range = 
  6 metrics
    VALUBusy, VALUUtilization, SQ_INSTS_MFMA, FetchSize, WriteSize, SQ_WAVES

ROCPRofiler: 1 contexts collected, output directory /tmp/rpl_data_250423_052539_921/input0_results_250423_052539
File 'rocprof_metrics_group2.csv' is generating

--- rocprof compute_throughput execution stderr ---

--- rocprof compute_throughput exit code: 0 ---
--- Contents of rocprof_metrics_group2.csv ---
Index,KernelName,gpu-id,queue-id,queue-index,pid,tid,grd,wgr,lds,scr,arch_vgpr,accum_vgpr,sgpr,wave_size,sig,obj,VALUBusy,VALUUtilization,SQ_INSTS_MFMA,FetchSize,WriteSize,SQ_WAVES,DispatchNs,BeginNs,EndNs,CompleteNs
0,"_gemm_fp8_blockscale_kernel.kd",2,0,0,985,985,98304,512,12288,0,112,0,48,64,0x0,0x7fd4a88e4740,1.2338271940,99.9625214163,688128,57924.6250000000,3072.0000000000,1536,994562653402370,994562783754982,994562783878983,994562783888597

Restored ROCm clock and settings to default with 'rocm-smi -r'.
--- Cleaning up temporary profiling files ---
--- End profiling attempt ---
HARDCODED_CONFIG Used (SPECIFIC for M=1024, N=1536, K=7168): {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=1024, N=3072, K=1536): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=1024, N=576, K=7168): {'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=1024, N=7168, K=256): {'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=1024, N=7168, K=2048): {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=1024, N=4608, K=7168): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=1024, N=7168, K=2304): {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=1024, N=512, K=7168): {'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=1024, N=4096, K=512): {'BLOCK_SIZE_M': 32, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=6144, N=1536, K=7168): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=6144, N=3072, K=1536): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=6144, N=576, K=7168): {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=6144, N=7168, K=256): {'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 16, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=6144, N=7168, K=2048): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=6144, N=4608, K=7168): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=6144, N=7168, K=2304): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=6144, N=512, K=7168): {'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 4, 'num_stages': 2}
HARDCODED_CONFIG Used (SPECIFIC for M=6144, N=4096, K=512): {'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'num_warps': 8, 'num_stages': 2}
