benchmarks/kernelbench/evaluator/evaluator.py

"""
Evaluator for KernelBench problems using kernelbench evaluation logic.

This evaluator can run inside a Docker container or as a native Python script,
and evaluates candidate kernel programs against KernelBench reference implementations.
"""

import os
import re
import subprocess
import sys
import tempfile
import traceback
from pathlib import Path


def ensure_run_and_check(evaluator_dir: Path):
    """Download run_and_check.py if not present.

    This allows the evaluator to work in native Python mode without Docker,
    automatically fetching the KernelBench evaluation script on first use.

    Args:
        evaluator_dir: Directory where the evaluator is located

    Returns:
        Path to run_and_check.py
    """
    run_and_check_path = evaluator_dir / "run_and_check.py"

    if not run_and_check_path.exists():
        import urllib.request

        commit = "423217d"
        url = f"https://raw.githubusercontent.com/ScalingIntelligence/KernelBench/{commit}/scripts/run_and_check.py"

        print(
            f"[INFO] Downloading run_and_check.py from KernelBench (commit {commit})...",
            file=sys.stderr,
        )
        try:
            urllib.request.urlretrieve(url, run_and_check_path)
            print(f"[INFO] Downloaded to {run_and_check_path}", file=sys.stderr)
        except Exception as e:
            raise RuntimeError(f"Failed to download run_and_check.py: {e}")

    return run_and_check_path


def evaluate(program_path: str):
    """
    Evaluate a candidate kernel program against the reference using run_and_check.py.

    Args:
        program_path: Path to the candidate program file

    Returns:
        Dictionary with combined_score (higher is better) and optional artifacts
    """
    try:
        # Read configuration from environment variables
        # These are injected by the benchmark setup
        level = int(os.environ.get("KERNELBENCH_LEVEL", "1"))
        problem_id = int(os.environ.get("KERNELBENCH_PROBLEM_ID", "1"))
        eval_mode = os.environ.get("KERNELBENCH_EVAL_MODE", "local")
        gpu = os.environ.get("KERNELBENCH_GPU", "H100")
        num_correct_trials = int(os.environ.get("KERNELBENCH_NUM_CORRECT_TRIALS", "5"))
        num_perf_trials = int(os.environ.get("KERNELBENCH_NUM_PERF_TRIALS", "100"))
        timeout = int(os.environ.get("KERNELBENCH_TIMEOUT", "300"))

        # Read the program and wrap it in ModelNew class for KernelBench format
        with open(program_path, "r") as f:
            program_content = f.read()

        is_triton = bool(
            re.search(r"^(import triton|from triton)", program_content, flags=re.MULTILINE)
        )

        # Create a temporary file with ModelNew wrapper
        with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as tmp_file:
            # Replace class Model with class ModelNew (if not already ModelNew)
            converted_content = program_content
            if "class ModelNew" not in converted_content:
                converted_content = re.sub(
                    r"^class Model(?=[(:])", "class ModelNew", converted_content, flags=re.MULTILINE
                )
                # Fix super() calls - use modern Python 3 super() without arguments
                converted_content = re.sub(r"super\(Model,\s*self\)", "super()", converted_content)
                converted_content = re.sub(r"super\(Model,\s*cls\)", "super()", converted_content)

            tmp_file.write(converted_content)
            kernel_src_path = tmp_file.name

        try:
            # Ensure run_and_check.py is available (downloads if needed)
            evaluator_dir = Path(__file__).parent
            run_and_check_path = ensure_run_and_check(evaluator_dir)

            # Build command to run run_and_check.py
            cmd = [
                sys.executable,
                str(run_and_check_path),
                "ref_origin=kernelbench",
                f"level={level}",
                f"problem_id={problem_id}",
                f"kernel_src_path={kernel_src_path}",
                f"eval_mode={eval_mode}",
                f"gpu={gpu}",
                f"num_correct_trials={num_correct_trials}",
                f"num_perf_trials={num_perf_trials}",
                f"timeout={timeout}",
                "check_kernel=False",  # Disable static checker to allow reference code
            ]

            # Setting the backend is important for KernelBench triton evaluation to work
            if is_triton:
                cmd.append("backend=triton")

            # Set up environment
            env = os.environ.copy()

            # Run the evaluation from the evaluator directory
            print(f"[INFO] Running evaluation command: {' '.join(cmd)}", file=sys.stderr)
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=timeout,
                cwd=str(evaluator_dir),
                env=env,
            )
        finally:
            # Clean up temporary file
            try:
                os.unlink(kernel_src_path)
            except Exception:
                pass

        # Parse the output to extract speedup
        stdout = result.stdout
        stderr = result.stderr

        if result.returncode != 0:
            print(
                f"[ERROR] Evaluation failed with return code {result.returncode}", file=sys.stderr
            )
            print(f"[ERROR] stdout: {stdout}", file=sys.stderr)
            print(f"[ERROR] stderr: {stderr}", file=sys.stderr)
            return {
                "combined_score": -100.0,
                "error": f"Evaluation subprocess failed: {stderr[:500]}",
                "return_code": result.returncode,
            }

        # Extract speedup from output
        speedup_eager = None
        speedup_compile = None
        kernel_time = None
        ref_eager_time = None

        for line in stdout.split("\n"):
            if "Speedup over eager:" in line:
                match = re.search(r"([0-9.]+)x", line)
                if match:
                    speedup_eager = float(match.group(1))
            elif "Speedup over torch.compile:" in line:
                match = re.search(r"([0-9.]+)x", line)
                if match:
                    speedup_compile = float(match.group(1))
            elif "Custom Kernel exec time:" in line:
                match = re.search(r"([0-9.]+) ms", line)
                if match:
                    kernel_time = float(match.group(1))
            elif "PyTorch Reference Eager exec time:" in line:
                match = re.search(r"([0-9.]+) ms", line)
                if match:
                    ref_eager_time = float(match.group(1))

        # If we found speedup, use it as the score
        if speedup_eager is not None and speedup_eager > 0:
            return {
                "combined_score": float(speedup_eager),
                "speedup_over_eager": speedup_eager,
                "speedup_over_compile": speedup_compile,
                "kernel_time_ms": kernel_time,
                "ref_eager_time_ms": ref_eager_time,
                "eval_mode": eval_mode,
                "gpu": gpu,
            }
        else:
            # Kernel failed correctness or didn't compile
            # Extract only relevant output starting from [Eval]
            stdout_excerpt = stdout
            if "[Eval]" in stdout:
                eval_start = stdout.find("[Eval]")
                stdout_excerpt = stdout[eval_start:]

            # Take last 5000 chars if too long
            if len(stdout_excerpt) > 5000:
                stdout_excerpt = stdout_excerpt[-5000:]

            return {
                "combined_score": -100.0,
                "error": "Kernel failed correctness check or did not compile",
                "stdout_excerpt": stdout_excerpt,
            }

    except subprocess.TimeoutExpired:
        return {
            "combined_score": -1.0,
            "error": f"Evaluation timed out after {timeout} seconds",
        }
    except Exception as e:
        traceback.print_exc()
        return {
            "combined_score": -100.0,
            "error": f"Error during evaluation: {str(e)}",
            "error_type": type(e).__name__,
        }


if __name__ == "__main__":
    # Backwards-compat: bridges old evaluate() -> dict to the container JSON
    # protocol.  wrapper.py is copied from skydiscover/evaluation/wrapper.py.
    from wrapper import run

    run(evaluate)