# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

import os
import sys
import shutil
import tempfile
import argparse
import torch
import pandas as pd

from abc import abstractmethod
from aiter import logger
from operator import itemgetter
import time
from aiter import dtypes
from aiter.jit.utils.chip_info import get_gfx_runtime as _chip_get_gfx

INVALID_TIME = -1


def _read_csv(filepath, **kwargs):
    """Read CSV with automatic cleanup of common formatting issues:
    trailing tabs/spaces, extra unnamed columns, whitespace in headers/values.
    """
    df = pd.read_csv(filepath, **kwargs)
    df.columns = df.columns.str.strip()
    df = df.loc[:, ~df.columns.str.startswith("Unnamed:")]
    str_cols = df.select_dtypes(include=["object"]).columns
    for col in str_cols:
        df[col] = df[col].apply(lambda v: v.strip() if isinstance(v, str) else v)
    df.dropna(how="all", inplace=True)
    return df


class TunerCommon:
    ARG_DEFAULTS = {
        "verbose": False,
        "tune_file": "",
        "untune_file": "",
        "errRatio": 0.05,
        "batch": 100,
        "profile_file": "",  # for all results
        "timeout": None,  # 100s timeout for per test
        "warmup": 5,  # 5 warmup iters for profiling
        "iters": 101,  # 101 run iters for profiling
        "min_improvement_pct": 3.0,  # only write shapes improved by >= N%
    }
    dtype2bpe_dict = {
        dtypes.fp16: 2,
        dtypes.bf16: 2,
        dtypes.i16: 2,
        dtypes.fp8: 1,
        dtypes.fp8_e8m0: 1,
        dtypes.i8: 1,
        dtypes.i32: 4,
        dtypes.i4x2: 1,
        dtypes.fp4x2: 1,
        torch.uint8: 1,
        torch.uint32: 4,
        dtypes.fp32: 4,
        torch.int4: 1 / 2,
        torch.float8_e4m3fnuz: 1,
        torch.float8_e4m3fn: 1,
    }
    INVALID_TIME = -1  # op not support or error

    INF_TIME = float("inf")  # op time is too large
    INVLAID_ERR_RATIO = 1.0  # err ratio is too large

    def __init__(self, name, key, resultList, description=None):
        self.parser = argparse.ArgumentParser(description=description)
        self._setup_common_arguments()
        self._setup_specific_arguments()
        self.columns = key + resultList
        self.keys = key
        self.tunedf = None
        self.untunedf = None
        self.name = name
        self.topk = 1
        self.success = pd.DataFrame(columns=self.columns)
        self.failed = pd.DataFrame(columns=self.columns)

        self.remain_untuned = pd.DataFrame(columns=self.keys)
        self.sort_keys = key
        self.start_time = 0
        self.num_warmup = 10
        self.num_iters = 101

    def get_arg_defaults(self):
        """get default arguments"""
        return self.ARG_DEFAULTS.copy()

    def get_bpe(self, dtype):
        return self.dtype2bpe_dict[dtype]

    def set_run_iters(self, input, indtype):
        """set warm iters and run iter for profiling"""
        """suggest warm iters * time1_per_iter > 100us"""

    def _setup_common_arguments(self):
        """set common arguments"""
        defaults = self.get_arg_defaults()
        self.parser.add_argument(
            "--verbose", "-v", action="store_true", help="more info"
        )
        self.parser.add_argument(
            "-i",
            "--untune_file",
            default=defaults["untune_file"],
            dest="untune_file",
            required=False,
            help="input",
        )
        self.parser.add_argument(
            "-o",
            "--tune_file",
            default=defaults["tune_file"],
            dest="tune_file",
            required=False,
            help="output: tuning result store this file",
        )
        self.parser.add_argument(
            "--mp",
            type=int,
            default=torch.cuda.device_count(),
            help="Tuning on multiple GPUs using multiple processes",
        )
        self.parser.add_argument(
            "-k",
            "--splitK",
            action="store_true",
            required=False,
            help="Use splitK kernels",
        )
        self.parser.add_argument(
            "--shape_grouped",
            action="store_true",
            default=False,
            required=False,
            help="Group all kernel candidates for the same shape onto one GPU "
            "to eliminate cross-GPU timing variance (also saves generate_data calls)",
        )
        self.parser.add_argument(
            "--sort",
            type=dtypes.str2bool,
            default=defaults.get("sort", False),
            required=False,
            help="Arranged according to the keys (True/False)",
        )
        self.parser.add_argument(
            "--errRatio",
            type=float,
            default=defaults["errRatio"],
            help="tolerable error ratio, default 0.05.",
        )
        self.parser.add_argument(
            "--batch",
            type=int,
            default=defaults["batch"],
            help="split untuned shapes to batches to tune",
        )
        self.parser.add_argument(
            "--all",
            action="store_true",
            required=False,
            help="retune all shapes in tune_file if tune file and untune file are the same, or retune shapes in untune file if tune file and untune file are different",
        )
        self.parser.add_argument(
            "-o2",
            "--profile_file",
            default=defaults["profile_file"],
            required=False,
            help="output: all tuning results stored in this file",
        )
        self.parser.add_argument(
            "--warmup",
            type=int,
            default=defaults["warmup"],
            help="warmup iters for profiling",
        )
        self.parser.add_argument(
            "--iters",
            type=int,
            default=defaults["iters"],
            help="run iters for profiling",
        )
        self.parser.add_argument(
            "--timeout",
            type=int,
            default=defaults["timeout"],
            help="timeout for task group",
        )
        self.parser.add_argument(
            "--run_config",
            nargs="?",
            const=True,
            default=False,
            metavar="TUNED_CSV",
            help="Run production operator benchmark and exit (no tuning). "
            "If a tuned CSV path is given, read shapes and kernels from it; "
            "otherwise read shapes from -i and run with default kernels.",
        )
        self.parser.add_argument(
            "--compare",
            action="store_true",
            required=False,
            help="Run production-op benchmark before and after tuning, print compare results, and keep a compare candidate CSV.",
        )
        self.parser.add_argument(
            "--update_improved",
            action="store_true",
            required=False,
            help="With --compare, update the final tuned CSV for shapes improved by at least --min_improvement_pct, or when pre-run has no valid baseline but post-run passes.",
        )
        self.parser.add_argument(
            "--min_improvement_pct",
            dest="min_improvement_pct",
            type=float,
            default=defaults.get("min_improvement_pct", 3.0),
            help="With --compare --update_improved, update tuned CSV only when a valid pre/post benchmark shows at least this percent improvement. Shapes with no valid pre-run baseline but passing post-run are still allowed to update.",
        )

    def parse_args(self):
        args = self.parser.parse_args()
        if args.update_improved and not args.compare:
            self.parser.error("--update_improved requires --compare")
        return args

    @abstractmethod
    def _setup_specific_arguments(self):
        """set specific arguments"""
        pass

    @abstractmethod
    def pre_process(self, args):
        """pre_process tunedf and untunedf"""
        pass

    @abstractmethod
    def tune(self, untunedf, tunedf, args):
        """tune process, return all results"""
        pass

    @abstractmethod
    def getKernelName(self, kernel_id):
        """obtain name of the kernel from its id"""
        pass

    @abstractmethod
    def calculate(self, results, inbpe=2, outbpe=2):
        """calculate TFLOPS and bandwidth"""
        pass

    @abstractmethod
    def result_to_df(self, rets):
        """transfer results to dataframe"""
        pass

    def update_config_files(self, file_path: str, merge_name: str):
        path_list = file_path.split(os.pathsep) if file_path else []
        if len(path_list) <= 1:
            return file_path
        df_list = []
        ## merge config files
        ##example: AITER_CONFIG_GEMM_A4W4="/path1:/path2"

        df_list.append(_read_csv(path_list[0]))
        for i, path in enumerate(path_list[1:]):
            if os.path.exists(path):
                df = _read_csv(path)
                df_list.append(df)
            else:
                print(f"path {i+1}: {path} (not exist)")

        if len(df_list) > 1:
            all_cols = list(df_list[0].columns)
            for df in df_list[1:]:
                for c in df.columns:
                    if c not in all_cols:
                        insert_before = (
                            "tflops" if "tflops" in all_cols else all_cols[-1]
                        )
                        all_cols.insert(all_cols.index(insert_before), c)
            _FILL_DEFAULTS = {"xbf16": 0, "run_1stage": 0, "ksplit": 0}
            for j in range(len(df_list)):
                for c in all_cols:
                    if c not in df_list[j].columns:
                        df_list[j][c] = _FILL_DEFAULTS.get(c, 0)
                df_list[j] = df_list[j][all_cols]
        merge_df = pd.concat(df_list, ignore_index=True) if df_list else pd.DataFrame()
        dedup_keys = self.keys
        if "_tag" in merge_df.columns:
            merge_df["_tag"] = merge_df["_tag"].fillna("")
            dedup_keys = self.keys + ["_tag"]
        merge_df = (
            merge_df.sort_values("us")
            .drop_duplicates(subset=dedup_keys, keep="first")
            .reset_index(drop=True)
        )
        new_file_path = f"/tmp/{merge_name}.csv"
        merge_df.to_csv(new_file_path, index=False)
        return new_file_path

    def get_untuned_gemm_list(self, untuned_gemm_file):
        assert os.path.exists(
            untuned_gemm_file
        ), f"Not exist untuned file: {untuned_gemm_file}"
        untunedf = _read_csv(untuned_gemm_file)
        filtered_df = untunedf.drop_duplicates().reset_index(drop=True)
        return filtered_df

    def get_out_file(self, tuned_file):
        """if there are multiple tuned file, then write tuning result to the first file"""
        path_list = tuned_file.split(os.pathsep) if tuned_file else []
        assert path_list, "output tuned file is empty"
        return path_list[0]

    def get_tuned_gemm_list(self, tuned_gemm_file, columns=[]):
        all_tuned_file = self.update_config_files(tuned_gemm_file, self.name)
        if os.path.exists(all_tuned_file):
            try:
                column_order = _read_csv(all_tuned_file, nrows=0).columns.tolist()
                tunedf = _read_csv(all_tuned_file)
                tunedf = tunedf[column_order]
            except pd.errors.EmptyDataError:
                print(f"Empty tuned file: {all_tuned_file}")
                columns = self.columns if not columns else columns
                tunedf = pd.DataFrame(columns=columns)
        else:
            print(f"Not exist tuned file: {all_tuned_file}")
            columns = self.columns if not columns else columns
            tunedf = pd.DataFrame(columns=columns)
        return tunedf

    def get_retune_gemm_list(self, args):
        """get retune gemm list from tune_file and untune_file"""
        if args.untune_file is None:
            raise ValueError("untune_file must be specified for retuning")
        if self.get_out_file(args.tune_file) == args.untune_file:
            # retune all shapes in tune_file
            self.untunedf = self.get_untuned_gemm_list(args.untune_file)
            gfx = self.get_gfx()
            cu_num = self.get_cu_num()
            if "gfx" not in self.untunedf.columns:
                self.untunedf["gfx"] = gfx
            target_mask = (self.untunedf["gfx"] == gfx) & (
                self.untunedf["cu_num"] == cu_num
            )
            self.tunedf = self.untunedf[~target_mask]
            self.untunedf = self.untunedf[target_mask]
            self.untunedf = self.untunedf[self.keys]
        else:
            # retune shapes that are in both untune_file and tune_file
            untunedf = self.get_untuned_gemm_list(args.untune_file)
            gfx = self.get_gfx()
            cu_num = self.get_cu_num()
            if "cu_num" not in untunedf.columns:
                untunedf["gfx"] = gfx
                untunedf["cu_num"] = cu_num
            else:
                target_mask = untunedf["cu_num"] == cu_num
                if "gfx" in untunedf.columns:
                    target_mask = target_mask & (untunedf["gfx"] == gfx)
                else:
                    untunedf["gfx"] = gfx
                untunedf = untunedf[target_mask]
            self.untunedf = untunedf[self.keys]
            self.tunedf = self.get_tuned_gemm_list(args.tune_file)

            untunedf_cols = self.untunedf.columns
            mask = (
                self.tunedf[untunedf_cols]
                .apply(tuple, axis=1)
                .isin(self.untunedf[untunedf_cols].apply(tuple, axis=1))
            )
            if args.verbose:
                logger.info(f"retuning {mask.sum()} shapes")
                print(self.tunedf[mask])
            self.tunedf = self.tunedf[~mask]

    def update_tunedf(self, df_old, df_updates):
        """update tuned result to old df"""
        """ for shapes already tuned, we update the result inplace"""
        if df_updates.empty:
            return df_old
        key_columns = self.keys
        df_updates = df_updates.loc[:, self.columns]
        # Widen integer columns to object so that float/string updates don't
        # trigger a Pandas dtype-coercion error (e.g. tflops=0 stored as int64
        # cannot accept a float like 2.61).
        import numpy as np

        for col in df_old.columns:
            if col in df_updates.columns and df_old[col].dtype != df_updates[col].dtype:
                try:
                    common = np.result_type(df_old[col].dtype, df_updates[col].dtype)
                except TypeError:
                    common = object
                df_old[col] = df_old[col].astype(common)
        df_old["_tmp_key"] = df_old[key_columns].apply(tuple, axis=1)
        df_updates["_tmp_key"] = df_updates[key_columns].apply(tuple, axis=1)
        matched_keys = df_updates[df_updates["_tmp_key"].isin(df_old["_tmp_key"])][
            "_tmp_key"
        ].tolist()
        unmatched_keys = df_updates[~df_updates["_tmp_key"].isin(df_old["_tmp_key"])][
            "_tmp_key"
        ].tolist()
        for key in matched_keys:
            df_old.loc[df_old.index[df_old["_tmp_key"] == key][0]] = df_updates.loc[
                df_updates["_tmp_key"] == key
            ].values[0]
        if unmatched_keys:
            unmatched_rows = df_updates[
                df_updates["_tmp_key"].isin(unmatched_keys)
            ].copy()
            df_old = pd.concat([df_old, unmatched_rows], ignore_index=True)
        df_old.drop("_tmp_key", axis=1, inplace=True)
        df_updates.drop("_tmp_key", axis=1, inplace=True)
        return df_old

    def sortResults(self, tune_file, issorted, values):
        tunedf = _read_csv(tune_file)
        if issorted:
            tunedf = tunedf.sort_values(by=values)
        dedup_keys = self.keys
        if "_tag" in tunedf.columns:
            tunedf["_tag"] = tunedf["_tag"].fillna("")
            dedup_keys = self.keys + ["_tag"]
        tunedf = tunedf.drop_duplicates(
            subset=dedup_keys,
            keep="last",
        )
        tunedf.to_csv(tune_file, index=False)

    def get_cu_num(self):
        gpu = torch.cuda.current_device()
        device_properties = torch.cuda.get_device_properties(gpu)
        cu_num = device_properties.multi_processor_count
        return cu_num

    def get_gfx(self):
        return _chip_get_gfx()

    def post_process(self, rets, args, topk=-1, fast_mode=False):
        """post process, post process all results to return topk results"""
        rets = list(rets)
        if args.profile_file != "":
            if args.verbose:
                logger.info(f"saving profile to {args.profile_file}")
            profiledf = self.result_to_df(sorted(rets, key=itemgetter(0)))
            if os.path.exists(args.profile_file):
                old_df = _read_csv(args.profile_file)
            else:
                old_df = pd.DataFrame(columns=self.columns)
            profiledf = pd.concat([old_df, profiledf], ignore_index=True)
            profiledf.to_csv(args.profile_file, index=False, na_rep="Null")

        if fast_mode or topk == -1:
            return rets
        from collections import defaultdict

        grouped_rets = defaultdict(list)
        bestConfigs = []

        for info, us, max_err_ratio in rets:
            grouped_rets[info[0]].append((info[1:], us, max_err_ratio))

        grouped_results = list(grouped_rets.items())

        for info_key, time_list in grouped_results:
            tol_err_ratio = args.errRatio
            sorted_time = sorted(time_list, key=lambda x: x[1])
            filtered_time = [
                (info_ex, round(us, 4), max_err_ratio)
                for info_ex, us, max_err_ratio in sorted_time
                if max_err_ratio <= tol_err_ratio
                and us != self.INVALID_TIME
                and us != self.INF_TIME
            ]
            if len(filtered_time) == 0:
                logger.error(
                    f"error: no valid candidate found for {info_key}, please check the result or errRatio in all result file running with --profile_file"
                )

            effective_topk = min(topk, len(filtered_time))
            if effective_topk < topk:
                print(f"choose {effective_topk} kernels")
            self.topk = effective_topk
            best_config = [
                ((info_key, *info_ex), us, max_err_ratio)
                for info_ex, us, max_err_ratio in filtered_time[0:effective_topk]
            ]
            if not best_config:
                logger.info(f"No kernel can be used for {info_key}")
                best_config = [((info_key, *sorted_time[0][0]), self.INVALID_TIME, 1.0)]
            bestConfigs.extend(best_config)
        resultdf = self.result_to_df(bestConfigs)
        return resultdf

    def tune_summary(self, status):
        """Summary of tuning results"""
        logger.info("============= Tuning results Summary: ==============")
        tuning_time = round(time.time() - self.tune_start_time, 4)
        tunedf = pd.concat([self.success, self.failed])
        logger.info(
            f"Tuning {status}. tune {len(tunedf)} shapes, total tuning time is {tuning_time} seconds"
        )
        logger.info("Successfully tuned shapes:")
        if not self.success.empty:
            print(self.success, flush=True)
        logger.info("Failed shapes:")
        print(self.failed, flush=True)

        tunedf_subset = tunedf[self.untunedf.columns].astype(self.untunedf.dtypes)
        mask = self.untunedf.apply(tuple, axis=1).isin(
            tunedf_subset.apply(tuple, axis=1)
        )
        self.remain_untuned = self.untunedf[~mask]

        if not self.remain_untuned.empty:
            logger.info("untuned shapes:")
            print(self.remain_untuned)
        if not self.remain_untuned.empty or not self.failed.empty:
            logger.error(
                "\033[91m[Tuning not Finished]\033[0m some shapes are not tuned or all failed, please check the result file or tune with --profile_file to get more details"
            )
            sys.exit(1)

    @abstractmethod
    def result_to_csv(self, results, file, concat=False):
        """write result to csv file, all means concat all results to file"""
        pass

    def update_tflops_bw(self, tune_file):
        """update tflops and bw from old tune_file"""
        pass

    def run_config(self, args):
        """Run the production operator for each shape in the untuned CSV.
        Subclasses should override this to call the actual production operator.
        Returns a list of dicts: [{"shape": str, "us": float, "status": "ok"/"error"}]
        """
        logger.info(f"run_config not implemented for {self.name}, skipping benchmark")
        return []

    def _clear_op_caches(self):
        """Clear operator-specific config caches. Subclasses should override this
        to clear only their own caches."""
        pass

    def _set_config_env_for_run_config(self, args, config_file=None):
        """Set the config env var to point to a tuned config file, clear caches,
        and enable AITER_REBUILD so that run_config rebuilds with new configs.
        *config_file* overrides the default (``-o`` / ``args.tune_file``).
        """
        defaults = self.get_arg_defaults()
        env_name = defaults.get("config_env_name")
        if not env_name:
            # Must return a 2-tuple: callers always unpack into old_val, old_rebuild.
            return None, None
        output_file = config_file if config_file else self.get_out_file(args.tune_file)
        old_val = os.environ.get(env_name)
        os.environ[env_name] = output_file
        logger.info(f"Setting {env_name}={output_file} for benchmark")
        # Clear operator-specific config caches
        self._clear_op_caches()
        # Enable AITER_REBUILD (level 2: rm .so only, keep build cache for faster rebuild)
        # and clear module caches so operators rebuild with new config
        from aiter.jit import core as jit_core

        old_rebuild = jit_core.AITER_REBUILD
        jit_core.AITER_REBUILD = 2
        jit_core.get_module.cache_clear()
        # Reset rebuilded_list so all modules get rebuilt on next call
        jit_core.rebuilded_list = ["module_aiter_enum"]
        # Clear loaded modules dict (use getattr to avoid Python name mangling of __ prefix in class methods)
        mds = getattr(jit_core, "__mds", None)
        if mds is not None:
            mds.clear()
        # Clear get_config_file lru_cache so it re-reads the env var
        jit_core.AITER_CONFIGS.get_config_file.cache_clear()
        return old_val, old_rebuild

    def _restore_config_env(self, env_name, old_val, old_rebuild=0):
        """Restore the config env var and AITER_REBUILD to original values."""
        if env_name is None:
            return
        if old_val is None:
            os.environ.pop(env_name, None)
        else:
            os.environ[env_name] = old_val
        try:
            from aiter.jit import core as jit_core

            jit_core.AITER_REBUILD = old_rebuild
        except ImportError:
            pass

    def _emit_report_lines(self, lines, report_file=None):
        if report_file:
            with open(report_file, "a") as f:
                f.write("\n".join(lines) + "\n")
            return
        for line in lines:
            print(line, flush=True)

    def _split_benchmark_status(self, status):
        status = "" if status is None else str(status)
        if status == "ok":
            return "OK", ""
        if status.startswith("error:"):
            return "ERROR", status[len("error:") :].strip()
        if status.startswith("mismatch"):
            detail = status[len("mismatch") :].lstrip(":").strip()
            return "MISMATCH", detail or "output mismatch vs reference"
        if not status:
            return "UNKNOWN", ""
        return status.upper(), ""

    def _format_benchmark_keys(self, row):
        parts = []
        for key in self.keys:
            value = row.get(key, "")
            parts.append(f"{key}={value}")
        return "keys: " + ", ".join(parts)

    def _emit_repro_csv(self, failed_repros, report_file=None):
        """Emit a copy-pasteable CSV block for reproducing failed shapes."""
        if not failed_repros:
            return
        untuned_keys = [k for k in self.keys if k != "cu_num"]
        csv_header = ",".join(untuned_keys)
        lines = [
            "",
            f"============= Repro CSV ({len(failed_repros)} failed shapes) =============",
            "Copy the lines below into a CSV file to reproduce:",
            csv_header,
        ]
        for kd in failed_repros:
            lines.append(",".join(str(kd.get(k, "")) for k in untuned_keys))
        self._emit_report_lines(lines, report_file)

    def _get_benchmark_e2e_us(self, row, suffix=""):
        return getattr(row, f"benchmark_e2e_us{suffix}", -1)

    def _get_benchmark_kernel_us(self, row, suffix=""):
        return getattr(row, f"benchmark_kernel_us{suffix}", None)

    def _print_benchmark_results(
        self, label, results, report_file=None, shapes_df=None
    ):
        """Print benchmark results to stdout or append them to a report file."""
        if not results:
            self._emit_report_lines([f"{label}: no results"], report_file)
            return
        results_df = self._benchmark_results_to_df(results, shapes_df=shapes_df)
        lines = [f"============= {label} Benchmark Results ============="]
        has_kernel_us = (
            not results_df.empty
            and "benchmark_kernel_us" in results_df.columns
            and results_df["benchmark_kernel_us"].notna().any()
        )
        if has_kernel_us:
            header = (
                f"{'Shape':<40} | {'Kernel(us)':>10} | {'E2E(us)':>10} | {'Status':>8}"
            )
        else:
            header = f"{'Shape':<40} | {'E2E(us)':>10} | {'Status':>8}"
        lines.append(header)
        lines.append("-" * len(header))
        failed_repros = []
        if results_df.empty:
            for idx, r in enumerate(results):
                shape_str = r.get("shape", "unknown")
                e2e_us = r.get("e2e_us", -1)
                status = r.get("status", "unknown")
                status_summary, status_detail = self._split_benchmark_status(status)
                e2e_str = f"{e2e_us:.2f}" if e2e_us > 0 else "N/A"
                lines.append(f"{shape_str:<40} | {e2e_str:>10} | {status_summary:>8}")
                if status_detail:
                    lines.append(f"reason: {status_detail}")
                if status_summary in ("ERROR", "MISMATCH"):
                    if shapes_df is not None and idx < len(shapes_df):
                        key_dict = {
                            k: shapes_df.iloc[idx].get(k, "") for k in self.keys
                        }
                    elif self.untunedf is not None and idx < len(self.untunedf):
                        key_dict = {
                            k: self.untunedf.iloc[idx].get(k, "") for k in self.keys
                        }
                    else:
                        key_dict = {}
                    if key_dict:
                        failed_repros.append(key_dict)
            self._emit_report_lines(lines, report_file)
            self._emit_repro_csv(failed_repros, report_file)
            return
        for row in results_df.itertuples(index=False):
            shape_str = getattr(row, "shape", "unknown")
            e2e_us = self._get_benchmark_e2e_us(row)
            kernel_us = self._get_benchmark_kernel_us(row)
            status = getattr(row, "benchmark_status", "unknown")
            status_summary, status_detail = self._split_benchmark_status(status)
            e2e_str = f"{e2e_us:.2f}" if e2e_us > 0 else "N/A"
            if has_kernel_us:
                kernel_str = (
                    f"{kernel_us:.2f}"
                    if kernel_us is not None and pd.notna(kernel_us) and kernel_us > 0
                    else "N/A"
                )
                lines.append(
                    f"{shape_str:<40} | {kernel_str:>10} | {e2e_str:>10} | {status_summary:>8}"
                )
            else:
                lines.append(f"{shape_str:<40} | {e2e_str:>10} | {status_summary:>8}")
            key_dict = {key: getattr(row, key, "") for key in self.keys}
            lines.append(self._format_benchmark_keys(key_dict))
            if status_detail:
                lines.append(f"reason: {status_detail}")
            if status_summary in ("ERROR", "MISMATCH"):
                failed_repros.append(key_dict)
        self._emit_report_lines(lines, report_file)
        self._emit_repro_csv(failed_repros, report_file)

    def _print_comparison(self, pre_results, post_results, report_file=None):
        """Print comparison to stdout or append it to a report file."""
        if not pre_results or not post_results:
            self._emit_report_lines(
                ["Cannot print comparison: missing pre or post results"],
                report_file,
            )
            return
        pre_df = self._benchmark_results_to_df(pre_results)
        post_df = self._benchmark_results_to_df(post_results)
        if pre_df.empty or post_df.empty:
            self._emit_report_lines(
                ["Cannot print comparison: missing comparable benchmark rows"],
                report_file,
            )
            return
        comparison_df = pre_df.merge(
            post_df,
            on=self.keys,
            how="outer",
            suffixes=("_pre", "_post"),
        )
        comparison_df["shape"] = comparison_df["shape_pre"]
        missing_shape_mask = comparison_df["shape"].isna() | (
            comparison_df["shape"] == ""
        )
        comparison_df.loc[missing_shape_mask, "shape"] = comparison_df.loc[
            missing_shape_mask, "shape_post"
        ]
        lines = ["============= Tune Performance Comparison ============="]
        header = f"{'Shape':<40} | {'Pre-E2E(us)':>13} | {'Post-E2E(us)':>14} | {'Speedup':>8} | {'Status':>8}"
        lines.append(header)
        lines.append("-" * len(header))
        compare_failed_repros = []
        for row in comparison_df.itertuples(index=False):
            shape = getattr(row, "shape", "unknown")
            pre_us = self._get_benchmark_e2e_us(row, "_pre")
            post_us = self._get_benchmark_e2e_us(row, "_post")
            post_status = getattr(row, "benchmark_status_post", "error")
            if pd.isna(post_status):
                pre_str = f"{pre_us:.2f}" if pd.notna(pre_us) and pre_us > 0 else "N/A"
                lines.append(
                    f"{shape:<40} | {pre_str:>13} | {'N/A':>14} | {'N/A':>8} | {'MISS':>8}",
                )
                lines.append(
                    self._format_benchmark_keys(
                        {key: getattr(row, key, "") for key in self.keys}
                    )
                )
                continue
            status_summary, status_detail = self._split_benchmark_status(post_status)
            if pre_us > 0 and post_us > 0:
                speedup = pre_us / post_us
                speedup_str = f"{speedup:.2f}x"
            else:
                speedup_str = "N/A"
            pre_str = f"{pre_us:.2f}" if pre_us > 0 else "N/A"
            post_str = f"{post_us:.2f}" if post_us > 0 else "N/A"
            lines.append(
                f"{shape:<40} | {pre_str:>13} | {post_str:>14} | {speedup_str:>8} | {status_summary:>8}"
            )
            key_dict = {key: getattr(row, key, "") for key in self.keys}
            lines.append(self._format_benchmark_keys(key_dict))
            if status_detail:
                lines.append(f"reason: {status_detail}")
            if status_summary in ("ERROR", "MISMATCH"):
                compare_failed_repros.append(key_dict)
        self._emit_report_lines(lines, report_file)
        self._emit_repro_csv(compare_failed_repros, report_file)

    def _benchmark_results_to_df(self, results, shapes_df=None):
        columns = self.keys + [
            "shape",
            "benchmark_status",
            "benchmark_kernel_us",
            "benchmark_e2e_us",
        ]
        if shapes_df is None:
            shapes_df = self.untunedf
        if shapes_df is None or len(shapes_df) == 0 or not results:
            return pd.DataFrame(columns=columns)

        shapes_df = shapes_df[self.keys].reset_index(drop=True)
        limit = min(len(shapes_df), len(results))
        if len(shapes_df) != len(results):
            logger.warning(
                f"benchmark results count mismatch in {self.name}: "
                f"{len(results)} results for {len(shapes_df)} shapes; matching by row order"
            )

        rows = []
        for idx in range(limit):
            bench = results[idx] or {}
            row = shapes_df.iloc[idx].to_dict()
            row["shape"] = bench.get("shape", "")
            row["benchmark_status"] = bench.get("status", "unknown")
            row["benchmark_kernel_us"] = bench.get("kernel_us", None)
            row["benchmark_e2e_us"] = bench.get("e2e_us", -1)
            rows.append(row)
        return pd.DataFrame(rows, columns=columns)

    def _build_compare_update_plan(
        self, pre_results, post_results, threshold_percent, shapes_df=None
    ):
        pre_df = self._benchmark_results_to_df(pre_results, shapes_df=shapes_df)
        post_df = self._benchmark_results_to_df(post_results, shapes_df=shapes_df)
        columns = self.keys + [
            "shape",
            "pre_us",
            "post_us",
            "pre_status",
            "post_status",
            "improvement_pct",
            "update",
            "update_reason",
        ]
        if pre_df.empty or post_df.empty:
            return pd.DataFrame(columns=columns)

        comparison = pre_df.merge(
            post_df,
            on=self.keys,
            how="outer",
            suffixes=("_pre", "_post"),
        )
        comparison["shape"] = comparison["shape_pre"]
        missing_shape_mask = comparison["shape"].isna() | (comparison["shape"] == "")
        comparison.loc[missing_shape_mask, "shape"] = comparison.loc[
            missing_shape_mask, "shape_post"
        ]
        comparison["pre_us"] = comparison["benchmark_e2e_us_pre"]
        comparison["post_us"] = comparison["benchmark_e2e_us_post"]
        comparison["pre_status"] = comparison["benchmark_status_pre"]
        comparison["post_status"] = comparison["benchmark_status_post"]

        valid = (
            (comparison["pre_status"] == "ok")
            & (comparison["post_status"] == "ok")
            & (comparison["pre_us"] > 0)
            & (comparison["post_us"] > 0)
        )
        no_baseline = (
            (comparison["post_status"] == "ok")
            & (comparison["post_us"] > 0)
            & ~((comparison["pre_status"] == "ok") & (comparison["pre_us"] > 0))
        )
        comparison["improvement_pct"] = (
            (comparison["pre_us"] - comparison["post_us"])
            / comparison["pre_us"]
            * 100.0
        )
        comparison.loc[~valid, "improvement_pct"] = float("nan")
        comparison["update_reason"] = "skip"
        comparison.loc[
            valid & (comparison["improvement_pct"] >= threshold_percent),
            "update_reason",
        ] = "threshold_met"
        comparison.loc[no_baseline, "update_reason"] = "no_baseline"
        comparison["update"] = comparison["update_reason"] != "skip"
        return comparison[columns]

    def _print_compare_update_plan(
        self,
        comparison,
        threshold_percent,
        tuned_file=None,
        report_file=None,
        apply_updates=True,
    ):
        if comparison is None or comparison.empty:
            self._emit_report_lines(
                ["Compare-gated CSV update skipped: no comparable benchmark rows"],
                report_file,
            )
            return

        # Count actions
        update_count = len(comparison[comparison["update_reason"] == "threshold_met"])
        no_baseline_count = len(
            comparison[comparison["update_reason"] == "no_baseline"]
        )
        skip_count = len(comparison[comparison["update_reason"] == "skip"])
        total = len(comparison)

        target_desc = tuned_file if tuned_file else "tuned csv"
        verb = "Updated" if apply_updates else "Would update"

        lines = [
            "============= Compare Report =============",
            f"Total shapes: {total} | {verb}: {update_count + no_baseline_count} "
            f"(improved: {update_count}, new: {no_baseline_count}) | Skipped: {skip_count}",
            f"Threshold: >= {threshold_percent:.1f}% improvement to update {target_desc}",
            "",
        ]

        # Updated shapes first
        if update_count + no_baseline_count > 0:
            lines.append(f"--- {verb} ({update_count + no_baseline_count} shapes) ---")
            header = f"{'Shape':<40} | {'Pre(us)':>10} | {'Post(us)':>10} | {'Improve':>9} | {'Action':>18}"
            lines.append(header)
            lines.append("-" * len(header))
            for row in comparison.itertuples(index=False):
                if row.update_reason == "skip":
                    continue
                pre_str = (
                    f"{row.pre_us:.2f}"
                    if pd.notna(row.pre_us) and row.pre_us > 0
                    else "N/A"
                )
                post_str = (
                    f"{row.post_us:.2f}"
                    if pd.notna(row.post_us) and row.post_us > 0
                    else "N/A"
                )
                improve_str = (
                    f"{row.improvement_pct:.2f}%"
                    if pd.notna(row.improvement_pct)
                    else "N/A"
                )
                action = "UPDATE" if row.update_reason == "threshold_met" else "NEW"
                lines.append(
                    f"{row.shape:<40} | {pre_str:>10} | {post_str:>10} | {improve_str:>9} | {action:>18}"
                )
            lines.append("")

        # Skipped shapes
        if skip_count > 0:
            lines.append(f"--- Skipped ({skip_count} shapes) ---")
            header = f"{'Shape':<40} | {'Pre(us)':>10} | {'Post(us)':>10} | {'Improve':>9} | {'Reason':>18}"
            lines.append(header)
            lines.append("-" * len(header))
            for row in comparison.itertuples(index=False):
                if row.update_reason != "skip":
                    continue
                pre_str = (
                    f"{row.pre_us:.2f}"
                    if pd.notna(row.pre_us) and row.pre_us > 0
                    else "N/A"
                )
                post_str = (
                    f"{row.post_us:.2f}"
                    if pd.notna(row.post_us) and row.post_us > 0
                    else "N/A"
                )
                improve_str = (
                    f"{row.improvement_pct:.2f}%"
                    if pd.notna(row.improvement_pct)
                    else "N/A"
                )
                pre_summary, _ = self._split_benchmark_status(row.pre_status)
                post_summary, post_detail = self._split_benchmark_status(
                    row.post_status
                )
                if post_summary in ("ERROR", "MISMATCH"):
                    reason = f"post-{post_summary.lower()}"
                elif (
                    pd.notna(row.improvement_pct)
                    and row.improvement_pct < threshold_percent
                ):
                    reason = f"< {threshold_percent:.1f}% improve"
                else:
                    reason = "no improvement"
                lines.append(
                    f"{row.shape:<40} | {pre_str:>10} | {post_str:>10} | {improve_str:>9} | {reason:>18}"
                )

        if not apply_updates:
            lines.append("")
            lines.append("Re-run with --update_improved to apply.")

        self._emit_report_lines(lines, report_file)

    def _merge_compare_filtered_results(self, base_file, candidate_file, comparison):
        old_df = self.get_tuned_gemm_list(base_file)
        if not os.path.exists(candidate_file):
            return old_df

        candidate_df = self.get_tuned_gemm_list(candidate_file)
        if comparison is None or comparison.empty:
            return old_df

        improved_keys = set(
            comparison.loc[comparison["update"], self.keys]
            .astype(str)
            .apply(tuple, axis=1)
            .tolist()
        )
        if not improved_keys:
            return old_df

        def key_mask(df):
            if df.empty:
                return pd.Series([], index=df.index, dtype=bool)
            return df[self.keys].astype(str).apply(tuple, axis=1).isin(improved_keys)

        kept_old = old_df[~key_mask(old_df)].copy()
        improved_rows = candidate_df[key_mask(candidate_df)].copy()
        merged = pd.concat([kept_old, improved_rows], ignore_index=True)
        dedup_keys = list(self.keys)
        if "_tag" in merged.columns:
            merged["_tag"] = merged["_tag"].fillna("")
            dedup_keys.append("_tag")
        merged = merged.drop_duplicates(subset=dedup_keys, keep="last").reset_index(
            drop=True
        )
        return merged

    def _run_config_for_shapes(self, args, shapes_df, config_file=None):
        original_untunedf = self.untunedf
        shapes_df = shapes_df.reset_index(drop=True)
        self.untunedf = shapes_df
        try:
            if config_file is None:
                return self.run_config(args)
            defaults = self.get_arg_defaults()
            env_name = defaults.get("config_env_name")
            old_val, old_rebuild = self._set_config_env_for_run_config(
                args, config_file=config_file
            )
            try:
                return self.run_config(args)
            finally:
                self._restore_config_env(env_name, old_val, old_rebuild)
        finally:
            self.untunedf = original_untunedf

    def _init_compare_report(self, args, output_file, batch_size, total_batches):
        if not args.compare or (total_batches <= 1 and len(self.untunedf) <= 30):
            return None

        compare_dir = os.path.join(tempfile.gettempdir(), "aiter_compare")
        os.makedirs(compare_dir, exist_ok=True)
        base_name = os.path.splitext(os.path.basename(output_file))[0]
        fd, compare_report_file = tempfile.mkstemp(
            prefix=f"{base_name}.",
            suffix=".compare.txt",
            dir=compare_dir,
            text=True,
        )
        os.chmod(compare_report_file, 0o600)
        with os.fdopen(fd, "w") as f:
            f.write(
                f"Compare report for {self.name}\n"
                f"Shapes: {len(self.untunedf)}\n"
                f"Batch size: {batch_size}\n"
                f"Total batches: {total_batches}\n\n"
            )
        print(f"Compare results will be written to {compare_report_file}", flush=True)
        return compare_report_file

    def _init_compare_candidate_file(self, args, output_file):
        if not args.compare:
            return None

        compare_dir = os.path.join(tempfile.gettempdir(), "aiter_compare")
        os.makedirs(compare_dir, exist_ok=True)
        base_name, candidate_ext = os.path.splitext(os.path.basename(output_file))
        pid = os.getpid()
        compare_candidate_file = os.path.join(
            compare_dir, f"{base_name}.{pid}.candidate{candidate_ext or '.csv'}"
        )
        if os.path.exists(output_file):
            shutil.copyfile(output_file, compare_candidate_file)
        elif os.path.exists(compare_candidate_file):
            os.remove(compare_candidate_file)
        print(
            f"Compare candidate CSV will be written to {compare_candidate_file}",
            flush=True,
        )
        return compare_candidate_file

    def _emit_compare_batch_header(self, header, report_file=None):
        print(header, flush=True)
        if report_file:
            self._emit_report_lines([header], report_file)

    def _run_compare_benchmark(
        self,
        args,
        batch,
        header,
        result_label,
        report_file=None,
        config_file=None,
        print_results=True,
    ):
        self._emit_compare_batch_header(header, report_file)
        results = self._run_config_for_shapes(args, batch, config_file=config_file)
        if print_results:
            self._print_benchmark_results(
                result_label, results, report_file=report_file
            )
        return results

    def _create_batch_compare_output_file(
        self,
        args,
        results,
        output_file,
        processed_batches,
        compare_candidate_file=None,
    ):
        fd, batch_compare_output_file = tempfile.mkstemp(
            prefix=f"{self.name}_compare_batch_{processed_batches}_",
            suffix=".csv",
        )
        os.close(fd)
        candidate_base_file = (
            compare_candidate_file
            if compare_candidate_file and os.path.exists(compare_candidate_file)
            else output_file
        )
        if os.path.exists(candidate_base_file):
            shutil.copyfile(candidate_base_file, batch_compare_output_file)
        else:
            pd.DataFrame(columns=self.columns).to_csv(
                batch_compare_output_file, index=False
            )
        self.result_to_csv(results, batch_compare_output_file, not args.all)
        if os.path.exists(batch_compare_output_file):
            self.sortResults(batch_compare_output_file, args.sort, self.sort_keys)
            if compare_candidate_file:
                shutil.copyfile(batch_compare_output_file, compare_candidate_file)
        return batch_compare_output_file

    def _apply_compare_batch_results(
        self,
        args,
        batch,
        results,
        batch_pre_tune_results,
        output_file,
        processed_batches,
        total_batches,
        compare_report_file=None,
        compare_candidate_file=None,
    ):
        batch_compare_output_file = self._create_batch_compare_output_file(
            args,
            results,
            output_file,
            processed_batches,
            compare_candidate_file=compare_candidate_file,
        )
        try:
            batch_header = f"=== Running post-tune benchmark (verification) for batch {processed_batches}/{total_batches} ==="
            batch_post_tune_results = self._run_compare_benchmark(
                args,
                batch,
                batch_header,
                "Post-tune",
                report_file=compare_report_file,
                config_file=batch_compare_output_file,
                print_results=args.verbose,
            )
            batch_compare_plan = self._build_compare_update_plan(
                batch_pre_tune_results,
                batch_post_tune_results,
                args.min_improvement_pct,
                shapes_df=batch,
            )
            if args.update_improved:
                final_df = self._merge_compare_filtered_results(
                    output_file,
                    batch_compare_output_file,
                    batch_compare_plan,
                )
                final_df.to_csv(output_file, index=False)
                if os.path.exists(output_file):
                    self.sortResults(output_file, args.sort, self.sort_keys)
                self.tunedf = self.get_tuned_gemm_list(output_file)
            return batch_post_tune_results, batch_compare_plan
        finally:
            if os.path.exists(batch_compare_output_file):
                os.remove(batch_compare_output_file)

    def _record_completed_compare_batch(
        self,
        completed_pre_tune_results,
        completed_post_tune_results,
        compare_plans,
        batch_pre_tune_results,
        batch_post_tune_results,
        batch_compare_plan,
    ):
        completed_pre_tune_results.extend(batch_pre_tune_results or [])
        completed_post_tune_results.extend(batch_post_tune_results or [])
        compare_plans.append(batch_compare_plan)

    def _print_compare_summary(
        self,
        completed_pre_tune_results,
        completed_post_tune_results,
        compare_plans,
        threshold_percent,
        tuned_file,
        report_file=None,
        apply_updates=True,
        candidate_file=None,
    ):
        if not completed_pre_tune_results:
            return

        self._print_comparison(
            completed_pre_tune_results,
            completed_post_tune_results,
            report_file=report_file,
        )
        combined_compare_plan = (
            pd.concat(compare_plans, ignore_index=True).reset_index(drop=True)
            if compare_plans
            else pd.DataFrame()
        )
        self._print_compare_update_plan(
            combined_compare_plan,
            threshold_percent,
            tuned_file=tuned_file,
            report_file=report_file,
            apply_updates=apply_updates,
        )
        extra_lines = []
        if candidate_file:
            extra_lines.append(f"Compare candidate CSV written to {candidate_file}")
        if not apply_updates:
            extra_lines.append(
                "Final tuned CSV was not updated. Re-run with --update_improved to apply improved shapes."
            )
        if extra_lines:
            self._emit_report_lines(extra_lines, report_file)
        if report_file:
            print(f"Compare results written to {report_file}", flush=True)

    #
    def run(self, args, fast_mode=False):
        """tuner run function"""
        self.pre_process(args)

        # Resolve --run_config: can be False, True (no file), or a file path string.
        # Strict semantics:
        #   --run_config <tuned_csv>  -> tuned kernels using that config file
        #   --run_config              -> default kernels (no config env override)
        run_config_file = args.run_config if isinstance(args.run_config, str) else None

        # --run_config with tuned file: load shapes from the tuned CSV.
        # --run_config without file: keep shapes from -i (pre_process), run default kernels.
        # --compare: always use untuned shapes from -i (pre_process).
        if args.run_config and run_config_file:
            tunedf = self.get_tuned_gemm_list(run_config_file)
            if not tunedf.empty and self.keys[0] in tunedf.columns:
                cu = self.get_cu_num()
                if "cu_num" in tunedf.columns:
                    tunedf = tunedf[tunedf["cu_num"] == cu]
                self.untunedf = tunedf.drop_duplicates(subset=self.keys).reset_index(
                    drop=True
                )

        print(self.untunedf)
        output_file = self.get_out_file(args.tune_file)
        if args.verbose:
            logger.info(f"args: {args}")

        # --run_config: only run benchmark and exit (no tuning)
        if args.run_config:
            if self.untunedf.empty:
                logger.info("No shapes to benchmark, nothing to run")
                return pd.DataFrame()
            if run_config_file:
                defaults = self.get_arg_defaults()
                env_name = defaults.get("config_env_name")
                old_val, old_rebuild = self._set_config_env_for_run_config(
                    args, config_file=run_config_file
                )
                try:
                    print(
                        "=== Running production operator benchmark (tuned) ===",
                        flush=True,
                    )
                    results = self.run_config(args)
                    self._print_benchmark_results("Benchmark (tuned)", results)
                finally:
                    self._restore_config_env(env_name, old_val, old_rebuild)
            else:
                print(
                    "=== Running production operator benchmark (default) ===",
                    flush=True,
                )
                results = self.run_config(args)
                self._print_benchmark_results("Benchmark (default)", results)
            return self.tunedf if self.tunedf is not None else pd.DataFrame()

        # Only include batches that fully completed compare+update in the final summary.
        completed_pre_tune_results = []
        completed_post_tune_results = []
        compare_plans = []

        if len(self.untunedf) == 0:
            # self.update_tflops_bw(args.tune_file)
            self.sortResults(output_file, args.sort, self.sort_keys)
            logger.info(
                f"no shapes to be tuned, skip tuning, tuned file is {args.tune_file}"
            )
            return self.tunedf if self.tunedf is not None else pd.DataFrame()
        batch_size = min(args.batch, len(self.untunedf))
        total_batches = (len(self.untunedf) + batch_size - 1) // batch_size
        compare_report_file = self._init_compare_report(
            args, output_file, batch_size, total_batches
        )
        compare_candidate_file = self._init_compare_candidate_file(args, output_file)
        if args.verbose:
            logger.info(
                f"total shapes to be tuned: {len(self.untunedf) }, total_batches: {total_batches}, batch_size: {batch_size}"
            )
            if args.compare and not args.update_improved:
                logger.info(
                    f"compare candidate results will be written to {compare_candidate_file}"
                )
            else:
                logger.info(f"results will be written to {output_file}")
        processed_batches = 0
        completed_batches = 0
        results = []
        topk = -1 if fast_mode else 1
        self.tune_start_time = time.time()
        tuning_status = "Finished"
        try:
            for i in range(0, len(self.untunedf), batch_size):
                batch = self.untunedf.iloc[i : i + batch_size].reset_index(drop=True)
                processed_batches += 1
                batch_pre_tune_results = None
                if args.compare:
                    batch_header = f"=== Running pre-tune benchmark (batch {processed_batches}/{total_batches}) ==="
                    batch_pre_tune_results = self._run_compare_benchmark(
                        args,
                        batch,
                        batch_header,
                        "Pre-tune",
                        report_file=compare_report_file,
                        print_results=args.verbose,
                    )
                all_results = self.tune(batch, self.tunedf, args)
                if all_results:
                    results = self.post_process(all_results, args, topk)
                    if args.compare:
                        batch_post_tune_results, batch_compare_plan = (
                            self._apply_compare_batch_results(
                                args,
                                batch,
                                results,
                                batch_pre_tune_results,
                                output_file,
                                processed_batches,
                                total_batches,
                                compare_report_file=compare_report_file,
                                compare_candidate_file=compare_candidate_file,
                            )
                        )
                        self._record_completed_compare_batch(
                            completed_pre_tune_results,
                            completed_post_tune_results,
                            compare_plans,
                            batch_pre_tune_results,
                            batch_post_tune_results,
                            batch_compare_plan,
                        )
                    else:
                        self.result_to_csv(results, output_file, not args.all)
                    completed_batches += 1
                    logger.info(
                        f"processed {completed_batches} batches of {total_batches}, Processing Status ====> {round(completed_batches / total_batches,2)*100:.1f}% tuned in {self.name}"
                    )
                else:
                    logger.info(
                        f"tune result is none or all shape is tuned in {args.tune_file}!"
                    )
            if os.path.exists(output_file):
                self.sortResults(output_file, args.sort, self.sort_keys)
        except KeyboardInterrupt:
            tuning_status = "Interrupted"
            logger.error(
                f"interrupted by user, tuning stopped, {completed_batches} batches processed"
            )
        except Exception as e:
            tuning_status = "Error"
            logger.error(
                f"error in batch {processed_batches} of {total_batches} after {completed_batches} completed batches: {str(e)}",
                exc_info=True,
            )
        finally:
            tune_exit = None
            summary_exc = None
            try:
                self.tune_summary(tuning_status)
            except SystemExit as e:
                tune_exit = e
            except Exception as e:
                summary_exc = e
                logger.error(
                    f"tune_summary failed (tuning may still have written results): {e}",
                    exc_info=True,
                )
            if args.compare:
                self._print_compare_summary(
                    completed_pre_tune_results,
                    completed_post_tune_results,
                    compare_plans,
                    args.min_improvement_pct,
                    output_file,
                    report_file=compare_report_file,
                    apply_updates=args.update_improved,
                    candidate_file=compare_candidate_file,
                )
            if tune_exit is not None:
                raise tune_exit
            if summary_exc is not None:
                raise summary_exc


class GemmCommonTuner(TunerCommon):

    ARG_DEFAULTS = {
        **TunerCommon.ARG_DEFAULTS,
        "sort": True,  # Enable sorting by default for GEMM tuners
    }

    def __init__(
        self,
        name,
        key=["gfx", "cu_num", "M", "N", "K"],
        resultList=[
            "kernelId",
            "splitK",
            "us",
            "kernelName",
            "tflops",
            "bw",
            "errRatio",
        ],
        description=None,
    ):
        super().__init__(name, key, resultList, description)
        # Swap M and N positions to ensure N comes before M
        self.sort_keys = list(key)
        m_idx = self.sort_keys.index("M")
        n_idx = self.sort_keys.index("N")
        self.sort_keys[m_idx], self.sort_keys[n_idx] = (
            self.sort_keys[n_idx],
            self.sort_keys[m_idx],
        )

    def pre_process(self, args):
        if args.all:
            self.get_retune_gemm_list(args)
        else:
            self.untunedf = self.get_untuned_gemm_list(args.untune_file)
            self.untunedf["gfx"] = self.get_gfx()
            self.untunedf["cu_num"] = self.get_cu_num()
            self.untunedf = self.untunedf[self.keys]
            self.tunedf = self.get_tuned_gemm_list(args.tune_file)

            untunedf_cols = self.untunedf.columns
            if len(self.tunedf) != 0:
                mask = self.untunedf.apply(tuple, axis=1).isin(
                    self.tunedf[untunedf_cols].apply(tuple, axis=1)
                )
                if args.verbose:
                    logger.info("skiped tuned shapes:")
                    print(self.untunedf[mask])
                self.untunedf = self.untunedf[~mask].reset_index(drop=True)

    def calculate(self, results, bpes=(2, 2, 2)):
        """calculate TFLOPS and bandwidth"""
        ### bpes: (inbpe, w_bpe, outbpe)
        ### gemm flops,bw
        info, time, err_ratio = results
        if time == -1:
            return 0, 0
        if len(info[0]) >= 5:  # gfx-aware key: (gfx, cu_num, m, n, k, ...)
            _gfx, cu_num, m, n, k, *rest = info[0]
        else:  # legacy subclass key: (cu_num, m, n, k, ...)
            cu_num, m, n, k, *rest = info[0]
        flop = m * n * k * 2
        tflops = round(flop / (time * 1000000), 2)
        lhs_bpe, rhs_bpe, out_bpe = bpes
        bw = round(
            (m * k * lhs_bpe + n * k * rhs_bpe + m * n * out_bpe) / (time * 1e-6) / 1e9,
            2,
        )
        return tflops, bw

    def result_to_df(self, results):
        resultdf = pd.DataFrame(columns=self.columns)
        for el in results:
            info, time, err_ratio = el
            keys, kernelId, splitK, kernelName = info
            # Resolve kernel name for both success and failure (profile CSV / debugging).
            # Treat missing/NA like "" so we always look up CK kernel names; otherwise NaN
            # would serialize as "Null" via na_rep in to_csv.
            need_lookup = kernelName == "" or pd.isna(kernelName)
            resolved = self.getKernelName(kernelId) if need_lookup else kernelName
            if resolved is None or pd.isna(resolved):
                kernelName = "None"
            else:
                kernelName = str(resolved)
            tflops, bw = self.calculate(el)
            key_dict = dict(zip(self.keys, keys))

            if len(results) == self.topk:
                print(
                    f"Tuning result for {str(key_dict).strip('{}')} is kernelId={kernelId} {kernelName} {splitK=}, {time}us, {err_ratio=}, {tflops=} TFLOPS, {bw=} GB/s"
                )
            key_dict.update(
                {
                    "kernelId": [kernelId],
                    "splitK": [splitK],
                    "us": [time],
                    "kernelName": [kernelName],
                    "errRatio": [err_ratio],
                    "tflops": [tflops],
                    "bw": [bw],
                }
            )
            temp = pd.DataFrame(key_dict)
            if resultdf.empty:
                resultdf = temp
            else:
                resultdf = pd.concat([resultdf, temp], ignore_index=True)
        return resultdf

    def result_to_csv(self, resultdf, file, concat=False):
        """post process of tuning results"""
        old_df = self.get_tuned_gemm_list(file)
        self.failed = pd.concat(
            [
                self.failed,
                resultdf[
                    (resultdf["us"] == self.INVALID_TIME)
                    | (resultdf["us"] == self.INF_TIME)
                ],
            ],
            ignore_index=True,
        )
        self.success = pd.concat(
            [
                self.success,
                resultdf[
                    (resultdf["us"] != self.INVALID_TIME)
                    & (resultdf["us"] != self.INF_TIME)
                ],
            ],
            ignore_index=True,
        )
        update_tunedf = resultdf[
            (resultdf["us"] != self.INVALID_TIME) & (resultdf["us"] != self.INF_TIME)
        ]  # self.success
        if not concat:
            resultdf = self.update_tunedf(old_df, update_tunedf)
        else:
            resultdf = pd.concat([old_df, update_tunedf], ignore_index=True)
        resultdf.to_csv(file, index=False, na_rep="Null")

    def update_tflops_bw(self, file):
        resultdf = self.get_tuned_gemm_list(file)
        for i in range(len(resultdf)):
            if len(resultdf.loc[i]) == 8:
                *keys, kernelId, splitK, us, kernelName = tuple(resultdf.loc[i])
            else:
                (
                    *keys,
                    kernelId,
                    splitK,
                    us,
                    kernelName,
                    tflops,
                    bw,
                    errRatio,
                ) = resultdf.iloc[i]
            errRatio = 0
            keys = tuple(keys)
            info = (keys, kernelId, splitK, ""), us, errRatio
            tflops, bw = self.calculate(info)
            resultdf.loc[i, "tflops"] = tflops
            resultdf.loc[i, "bw"] = bw
            resultdf.loc[i, "errRatio"] = 0
        resultdf.to_csv(file, index=False, na_rep="Null")

    def set_run_iters(self, input, inputdtype):
        if len(input) >= 5:  # gfx-aware key: (gfx, cu_num, m, n, k, ...)
            _gfx, cu_num, m, n, k, *rest = input
        else:  # legacy subclass key: (cu_num, m, n, k, ...)
            cu_num, m, n, k, *rest = input
        flops = m * n * k * 2
        if flops < 256 * 5120 * 256 * 2:
            self.num_warmup = 50
        elif flops <= 1024 * 5120 * 256 * 2:
            self.num_warmup = 30
