#!/usr/bin/env python3
"""Evaluation orchestrator for NVFP4 kernel implementations.
Runs correctness, performance, and profiling checks on kernel backends.
Uses cached reference outputs for decoupled, isolated execution.
For focused testing, use specialized modules:
    - kernel_utils.correctness: Correctness-only testing
    - kernel_utils.benchmark: Performance benchmarking
    - kernel_utils.profiling: Torch/NCU profiling
Usage:
    # Test implementation kernel (cached mode - decoupled from reference)
    python -m kernel_utils.evaluate --implementation /path/to/kernel.py
    # Test with custom reference
    python -m kernel_utils.evaluate --implementation /path/to/kernel.py --reference /path/to/reference.py
    # Test with problem_id reference
    python -m kernel_utils.evaluate --implementation /path/to/kernel.py --reference nvfp4_gemv_blackwell
"""
from __future__ import annotations
import logging
import sys
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol
import torch

logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger(__name__)
# evaluate.py is at wafer_utils/kernel_utils/evaluate.py
# Need to go up 2 levels to reach research/ to import wafer_utils
sys.path.insert(0, str(Path(__file__).parent.parent.parent))
# Correctness checking tolerances
CORRECTNESS_RTOL = 1e-3  # Relative tolerance for FP16 output comparison
CORRECTNESS_ATOL = 1e-3  # Absolute tolerance for FP16 output comparison
# Benchmarking constants
BENCHMARK_NUM_WARMUP = 5  # Number of warmup runs before benchmarking
BENCHMARK_NUM_RUNS = 20  # Number of benchmark runs for averaging
# Profiling constants
PROFILE_NUM_WARMUP = 5  # Number of warmup runs before profiling
PROFILE_NUM_RUNS = 1  # Number of profiling runs (profiling is expensive)
# Time conversion constants
MS_TO_MICROSECONDS = 1000  # Convert milliseconds to microseconds
# Performance approximation factors (for display when detailed stats unavailable)
PERF_MIN_FACTOR = 0.99  # Minimum time approximation factor
PERF_MAX_FACTOR = 1.01  # Maximum time approximation factor
PERF_STD_FACTOR = 0.001  # Standard deviation approximation factor
# Error sentinel values
ERROR_TIME_MS = -1.0  # Sentinel value for failed benchmark times
# Display formatting
PRINT_SEPARATOR_WIDTH = 80  # Width of separator lines in output
from wafer_core.utils.exceptions import (  # noqa: E402
    InvalidTestCasesFormatError,
    ReferenceKernelLoadError,
)
from wafer_core.utils.kernel_utils.results import (  # noqa: E402
    BackendResults,
    CorrectnessResult,
    PerformanceResult,
    TestSuiteResults,
)
from wafer_core.utils.kernel_utils.task import (  # noqa: E402
    TestCase,
)
from wafer_core.utils.kernel_utils.utils import (  # noqa: E402
    ErrorInfo,
    benchmark_kernel,
    benchmark_kernel_defensive,
    make_match_reference,
    ncu_profile_kernel,
    profile_kernel,
)
_reference_kernel_cache: dict[str, Callable] = {}
# Parameter dataclasses for reducing function argument count
# Tiger Style: Hourglass shape - group related params together
@dataclass(frozen=True)
class TestExecutionContext:
    """Context for executing a single test (paths and configuration).
    Immutable context passed down the test execution stack.
    """
    implementation_path: str
    reference_path: str
    test_suite: str
    use_cached_reference: bool
@dataclass(frozen=True)
class ProfilingConfig:
    """Configuration for profiling operations.
    Groups profiling-related settings together.
    """
    profile_subdir: str
    ncu_subdir: str
    enable_profiling: bool
    enable_ncu: bool
    use_defenses: bool = False
@dataclass(frozen=True)
class TestRunConfig:
    """Complete configuration for a test run.
    Combines execution context and profiling config.
    """
    context: TestExecutionContext
    profiling: ProfilingConfig
def _load_reference_kernel(problem_id: str) -> Callable:
    """Dynamically load reference kernel based on problem_id.
    Args:
        problem_id: Problem identifier (e.g., "nvfp4_gemv_blackwell", "gpumode_540")
    Returns:
        Reference kernel solve function
    Raises:
        ReferenceKernelLoadError: If problem_id doesn't have a reference kernel module
    """
    assert problem_id, "problem_id cannot be empty"
    assert isinstance(problem_id, str), f"problem_id must be string, got {type(problem_id)}"
    if problem_id in _reference_kernel_cache:
        return _reference_kernel_cache[problem_id]
    # Try to import reference kernel from problem module
    # Format: {problem_id}/reference_kernel.py -> solve function
    try:
        if problem_id.startswith("gpumode_"):
            # For GPUMode problems, try to import from a module matching the problem_id
            module_name = f"{problem_id}.reference_kernel"
        else:
            # Direct problem_id (e.g., "nvfp4_gemv_blackwell")
            module_name = f"{problem_id}.reference_kernel"
        module = __import__(module_name, fromlist=["solve"])
        reference_solve = module.solve
        _reference_kernel_cache[problem_id] = reference_solve
        logger.info(f"loaded reference kernel for problem: {problem_id}")
    except ImportError as e:
        # Fail fast - no fallbacks
        raise ReferenceKernelLoadError(problem_id, e) from e
    else:
        return reference_solve
def ref_kernel_adapter(inputs: Any, problem_id: str = "nvfp4_gemv_blackwell") -> Any:
    """Adapter to convert dict inputs to reference kernel.
    Args:
        inputs: Input data (dict or tuple)
        problem_id: Problem identifier for dynamic reference kernel loading
    """
    assert inputs is not None, "inputs cannot be None"
    assert problem_id, "problem_id cannot be empty"
    reference_solve = _load_reference_kernel(problem_id)
    if isinstance(inputs, dict):
        return reference_solve(**inputs)
    else:
        # Old tuple interface: (a, b, scale_a, scale_b, scale_a_permuted, scale_b_permuted, c)
        a, b, scale_a, scale_b, scale_a_permuted, scale_b_permuted, c = inputs
        return reference_solve(
            a=a,
            b=b,
            scale_a=scale_a,
            scale_b=scale_b,
            scale_a_cpu=scale_a.cpu(),
            scale_b_cpu=scale_b.cpu(),
            scale_a_permuted=scale_a_permuted,
            scale_b_permuted=scale_b_permuted,
            c=c,
        )
def load_reference_kernel_from_path(reference_path: str) -> Callable:
    """Load reference kernel from path or problem_id.
    Args:
        reference_path: Either a file path or problem_id string (e.g., "nvfp4_gemv_blackwell")
    Returns:
        Reference kernel function
    Raises:
        FileNotFoundError: If reference_path is a path but file doesn't exist
        ReferenceKernelLoadError: If reference_path is a problem_id but can't be loaded
    """
    assert reference_path, "reference_path cannot be empty"
    from pathlib import Path
    ref_path = Path(reference_path)
    if ref_path.exists() and ref_path.is_file():
        return load_kernel_from_file(reference_path, function_name="ref_kernel")
    elif "/" in reference_path or "\\" in reference_path or reference_path.endswith(".py"):
        # Looks like a file path but doesn't exist - fail fast
        raise FileNotFoundError(
            f"Reference file not found: {reference_path}. "
            "If this is a problem_id, it should not contain path separators or .py extension."
        )
    else:
        # Treat as problem_id and load built-in reference (will fail fast if not found)
        return ref_kernel_adapter
def _print_gpumode_correctness_results(backend_results: BackendResults) -> None:
    """Print correctness results in gpumode_test.txt format."""
    assert backend_results is not None, "backend_results cannot be None"
    print("\n📋 Correctness Summary:")
    for result in backend_results.correctness_tests:
        status = "✅" if result.is_correct else "❌"
        params = dict(item.split("=") for item in result.test_params.split(", "))
        print(
            f"{status} {result.test_name}: k={params['k']}, l={params['l']}, m={params['m']}, seed={params['seed']}"
        )
def _print_gpumode_benchmark_results(backend_results: BackendResults) -> None:
    """Print benchmark results in gpumode_benchmark.txt and leaderboard format."""
    import math
    assert backend_results is not None, "backend_results cannot be None"
    print("\n📊 gpumode_benchmark.txt:")
    benchmark_times = []
    for result in backend_results.performance_tests:
        if result.test_name.startswith("gpumode_bench") and result.successfully_ran:
            params = dict(item.split("=") for item in result.test_params.split(", "))
            avg_time_us = result.avg_time_ms * MS_TO_MICROSECONDS
            # We don't have min/max/std from current implementation, so approximate
            # In real implementation, we'd need to track these in PerformanceResult
            min_time_us = avg_time_us * PERF_MIN_FACTOR
            max_time_us = avg_time_us * PERF_MAX_FACTOR
            std_dev_us = avg_time_us * PERF_STD_FACTOR
            print(f"k: {params['k']}; l: {params['l']}; m: {params['m']}; seed: {params['seed']}")
            print(f" ⏱ {avg_time_us:.1f} ± {std_dev_us:.2f} µs")
            print(f" ⚡ {min_time_us:.1f} µs 🐌 {max_time_us:.1f} µs")
            print()
            benchmark_times.append(avg_time_us)
    # Print leaderboard with geometric mean score
    if benchmark_times:
        print("🏆 gpumode_leaderboard.txt:")
        for result in backend_results.performance_tests:
            if result.test_name.startswith("gpumode_bench") and result.successfully_ran:
                params = dict(item.split("=") for item in result.test_params.split(", "))
                avg_time_us = result.avg_time_ms * MS_TO_MICROSECONDS
                min_time_us = avg_time_us * PERF_MIN_FACTOR
                max_time_us = avg_time_us * PERF_MAX_FACTOR
                std_dev_us = avg_time_us * PERF_STD_FACTOR
                print(
                    f"k: {params['k']}; l: {params['l']}; m: {params['m']}; seed: {params['seed']}"
                )
                print(f" ⏱ {avg_time_us:.1f} ± {std_dev_us:.2f} µs")
                print(f" ⚡ {min_time_us:.1f} µs 🐌 {max_time_us:.1f} µs")
                print()
        geomean_score = math.exp(sum(math.log(t) for t in benchmark_times) / len(benchmark_times))
        print(f" score: {geomean_score:.3f}μs")
def print_gpumode_format(
    backend_results: BackendResults,
    test_suites: list[str],
    test_cases: list[TestCase],
) -> None:
    """Print results in GPUMode competition format.
    Args:
        backend_results: Results for a single backend
        test_suites: List of test suite names that were run
        test_cases: List of test cases that were run
    """
    assert backend_results is not None, "backend_results cannot be None"
    assert test_suites is not None, "test_suites cannot be None"
    # Only print GPUMode format if we ran GPUMode test suites
    if "gpumode_correctness" not in test_suites and "gpumode_benchmark" not in test_suites:
        return
    print("\n" + "=" * PRINT_SEPARATOR_WIDTH)
    print("GPUMode Competition Format")
    print("=" * PRINT_SEPARATOR_WIDTH)
    # Print correctness tests
    if "gpumode_correctness" in test_suites:
        _print_gpumode_correctness_results(backend_results)
    # Print benchmark tests
    if "gpumode_benchmark" in test_suites:
        _print_gpumode_benchmark_results(backend_results)
def _setup_test_backend(
    implementation_path: str,
    reference_path: str,
    test_suites: list[str],
) -> tuple[object, object | None, bool, str]:
    """Setup implementation and checker function.
    Returns:
        (implementation_fn, checker, use_cached_reference, error_msg): Implementation and checker on success, error on failure
    """
    assert implementation_path, "implementation_path cannot be empty"
    assert reference_path, "reference_path cannot be empty"
    assert test_suites is not None, "test_suites cannot be None"
    assert len(test_suites) > 0, "test_suites cannot be empty"
    try:
        implementation_fn = load_kernel_from_file(implementation_path)
    except Exception as e:
        return None, None, False, f"Failed to load implementation from {implementation_path}: {e}"
    # Always use live reference mode (no caching for now)
    # We explicitly pass both implementation and reference paths,
    # so we can run them both directly without pre-generated caches
    use_cached_reference = False
    print(f"   Using live reference: {reference_path}")
    try:
        ref_fn = load_reference_kernel_from_path(reference_path)
        checker = make_match_reference(ref_fn, rtol=CORRECTNESS_RTOL, atol=CORRECTNESS_ATOL)
    except Exception as e:
        return None, None, False, f"Failed to load reference from {reference_path}: {e}"
    return implementation_fn, checker, use_cached_reference, ""
def _run_correctness_test(
    test: TestCase,
    implementation_fn: Callable,
    checker: Callable,
    use_cached_reference: bool,
    test_suite: str,
) -> tuple[bool, ErrorInfo | None, object]:
    """Run correctness check on a single test.
    Returns:
        (is_correct, error_msg, test_input): Correctness result and input data.
        error_msg is ErrorInfo if test failed, None otherwise.
    """
    assert test is not None, "test cannot be None"
    assert implementation_fn is not None, "implementation_fn cannot be None"
    assert test_suite, "test_suite cannot be empty"
    assert checker is not None, "checker cannot be None"
    # Note: use_cached_reference is always False (we use live reference mode)
    # Keeping the parameter for now to avoid breaking the call chain
    test_input = test.generate()
    is_correct, error_msg = checker(implementation_fn, test_input)
    return is_correct, error_msg, test_input
def _should_stop_early(error_msg: ErrorInfo | None) -> bool:
    """Check if error indicates fundamental failure requiring early stop.
    Fundamental errors (coding bugs) vs numerical errors (tolerance issues).
    """
    if not error_msg:
        return False
    error_string = error_msg.to_string()
    fundamental_errors = [
        "TypeError",
        "ValueError",
        "AttributeError",
        "NameError",
        "SyntaxError",
        "ImportError",
        "CUDA error",  # Any CUDA error
        "AcceleratorError",
    ]
    return any(err in error_string for err in fundamental_errors)
def _run_performance_test(
    test: TestCase,
    implementation_fn: object,
    implementation_path: str,
    reference_path: str,
    test_input: object | None,
    use_cached_reference: bool,
    use_defenses: bool = False,
) -> PerformanceResult:
    """Run performance benchmark and calculate speedup.
    Args:
        test: Test case to run
        implementation_fn: Implementation kernel function
        implementation_path: Path to implementation file
        reference_path: Path to reference file
        test_input: Pre-generated test input (optional)
        use_cached_reference: Whether to use cached reference outputs
        use_defenses: Enable defensive timing to detect evaluation hacking
    Returns:
        PerformanceResult with benchmark times and speedup
    """
    assert test is not None, "test cannot be None"
    assert implementation_fn is not None, "implementation_fn cannot be None"
    assert implementation_path, "implementation_path cannot be empty"
    assert reference_path, "reference_path cannot be empty"
    # Tiger Style: Type narrow to Callable after assertion
    assert callable(implementation_fn), "implementation_fn must be callable"
    impl_fn: Callable = implementation_fn
    if use_cached_reference and test_input is not None:
        test_input_bench = test_input
    else:
        test_input_bench = test.generate()
    defense_results = None
    if use_defenses:
        print("   🛡️  Using defensive timing (detecting eval hacking)")
        avg_time, defense_results, bench_err = benchmark_kernel_defensive(
            impl_fn, test_input_bench, num_warmup=BENCHMARK_NUM_WARMUP, num_runs=BENCHMARK_NUM_RUNS
        )
    else:
        avg_time, bench_err = benchmark_kernel(
            impl_fn, test_input_bench, num_warmup=BENCHMARK_NUM_WARMUP, num_runs=BENCHMARK_NUM_RUNS
        )
    speedup = None
    ref_time = None
    if implementation_path != reference_path and bench_err is None:
        if use_cached_reference:
            # In cached mode, don't re-run reference for performance
            print("   ⚠️  Speedup calculation skipped (reference not benchmarked in cached mode)")
        else:
            ref_fn = load_reference_kernel_from_path(reference_path)
            test_input_ref = test.generate()
            ref_time, ref_err = benchmark_kernel(
                ref_fn, test_input_ref, num_warmup=BENCHMARK_NUM_WARMUP, num_runs=BENCHMARK_NUM_RUNS
            )
            if ref_err is None:
                speedup = ref_time / avg_time if avg_time > 0 else 0.0
    return PerformanceResult(
        test_name=test.name,
        backend_name=implementation_path,
        successfully_ran=(bench_err is None),
        test_params=test.serialize(),
        avg_time_ms=avg_time if bench_err is None else ERROR_TIME_MS,
        speedup=speedup,
        reference_time_ms=ref_time,
        error_msg=bench_err,
    )
def _run_torch_profiling(
    test: TestCase,
    implementation_fn: object,
    implementation_path: str,
    reference_path: str,
    profile_subdir: str,
) -> list:
    """Run torch profiling for implementation and reference.
    Returns:
        List of (implementation_path, test_name, trace_path) tuples
    """
    assert test is not None, "test cannot be None"
    assert implementation_fn is not None, "implementation_fn cannot be None"
    assert implementation_path, "implementation_path cannot be empty"
    # Tiger Style: Type narrow to Callable after assertion
    assert callable(implementation_fn), "implementation_fn must be callable"
    impl_fn: Callable = implementation_fn
    profile_traces = []
    test_input_prof = test.generate()
    profile_dir = Path(profile_subdir)
    trace_path, prof_err = profile_kernel(
        impl_fn,
        test_input_prof,
        profile_dir,
        implementation_path,
        test.name,
        num_warmup=PROFILE_NUM_WARMUP,
        num_profile_runs=PROFILE_NUM_RUNS,
    )
    if prof_err is None:
        print(f"   📊 Torch profile saved: {trace_path}")
        profile_traces.append((implementation_path, test.name, trace_path))
    else:
        print(f"   ⚠️  Torch profiling failed: {prof_err}")
    # Also profile reference for comparison
    if implementation_path != reference_path:
        ref_fn = load_reference_kernel_from_path(reference_path)
        test_input_ref_prof = test.generate()
        ref_trace_path, ref_prof_err = profile_kernel(
            ref_fn,
            test_input_ref_prof,
            profile_dir,
            reference_path,
            test.name,
            num_warmup=PROFILE_NUM_WARMUP,
            num_profile_runs=PROFILE_NUM_RUNS,
        )
        if ref_prof_err is None:
            print(f"   📊 Reference torch profile saved: {ref_trace_path}")
            profile_traces.append((reference_path, test.name, ref_trace_path))
        else:
            print(f"   ⚠️  Reference torch profiling failed: {ref_prof_err}")
    return profile_traces
def _run_ncu_profiling(
    test: TestCase,
    implementation_fn: object,
    implementation_path: str,
    reference_path: str,
    ncu_subdir: str,
) -> list:
    """Run NCU profiling for implementation and reference.
    Returns:
        List of (implementation_path, test_name, ncu_path) tuples
    """
    assert test is not None, "test cannot be None"
    assert implementation_fn is not None, "implementation_fn cannot be None"
    assert implementation_path, "implementation_path cannot be empty"
    # Tiger Style: Type narrow to Callable after assertion
    assert callable(implementation_fn), "implementation_fn must be callable"
    impl_fn: Callable = implementation_fn
    ncu_reports = []
    test_input_ncu = test.generate()
    ncu_dir = Path(ncu_subdir)
    # Profile implementation
    ncu_report_path, ncu_err = ncu_profile_kernel(
        impl_fn, test_input_ncu, ncu_dir, implementation_path, test.name, test
    )
    if ncu_err is None:
        print(f"   📊 NCU report saved: {ncu_report_path}")
        ncu_reports.append((implementation_path, test.name, ncu_report_path))
    else:
        print(f"   ⚠️  NCU profiling failed: {ncu_err}")
    # Also profile reference for comparison
    if implementation_path != reference_path:
        ref_fn = load_reference_kernel_from_path(reference_path)
        test_input_ref_ncu = test.generate()
        ref_ncu_path, ref_ncu_err = ncu_profile_kernel(
            ref_fn, test_input_ref_ncu, ncu_dir, reference_path, test.name, test
        )
        if ref_ncu_err is None:
            print(f"   📊 Reference NCU report saved: {ref_ncu_path}")
            ncu_reports.append((reference_path, test.name, ref_ncu_path))
        else:
            print(f"   ⚠️  Reference NCU profiling failed: {ref_ncu_err}")
    return ncu_reports
def _run_profiling(
    test: TestCase,
    implementation_fn: object,
    context: TestExecutionContext,
    profiling_config: ProfilingConfig,
) -> tuple[list, list]:
    """Run torch and NCU profiling if enabled.
    Returns:
        (profile_traces, ncu_reports): Lists of profiling outputs
    """
    assert test is not None, "test cannot be None"
    assert implementation_fn is not None, "implementation_fn cannot be None"
    profile_traces = []
    ncu_reports = []
    # Only profile benchmark tests
    is_benchmark_test = "bench" in test.name.lower() or "perf" in test.name.lower()
    if not is_benchmark_test:
        return profile_traces, ncu_reports
    # Torch profiling
    if profiling_config.enable_profiling:
        profile_traces = _run_torch_profiling(
            test,
            implementation_fn,
            context.implementation_path,
            context.reference_path,
            profiling_config.profile_subdir,
        )
    # NCU profiling
    if profiling_config.enable_ncu:
        ncu_reports = _run_ncu_profiling(
            test,
            implementation_fn,
            context.implementation_path,
            context.reference_path,
            profiling_config.ncu_subdir,
        )
    return profile_traces, ncu_reports
def _create_skipped_test_results(
    test: TestCase,
    implementation_path: str,
) -> tuple[CorrectnessResult, PerformanceResult]:
    """Create skipped test results for early-stop scenarios.
    Returns:
        (correctness_result, performance_result): Both marked as failed/skipped
    """
    assert test is not None, "test cannot be None"
    assert implementation_path, "implementation_path cannot be empty"
    correctness_result = CorrectnessResult(
        test_name=test.name,
        backend_name=implementation_path,
        is_correct=False,
        test_params=test.serialize(),
        error_msg="Skipped - kernel failed earlier tests",
    )
    performance_result = PerformanceResult(
        test_name=test.name,
        backend_name=implementation_path,
        successfully_ran=False,
        test_params=test.serialize(),
        avg_time_ms=ERROR_TIME_MS,
        speedup=None,
        reference_time_ms=None,
        error_msg="Skipped - kernel failed earlier tests",
    )
    return correctness_result, performance_result
def _run_performance_and_profiling(
    test: TestCase,
    implementation_fn: object,
    test_input: object,
    context: TestExecutionContext,
    profiling_config: ProfilingConfig,
) -> tuple[PerformanceResult, list, list]:
    """Run performance benchmark and profiling for a test.
    Returns:
        (performance_result, profile_traces, ncu_reports)
    """
    assert test is not None, "test cannot be None"
    assert implementation_fn is not None, "implementation_fn cannot be None"
    assert test_input is not None, "test_input cannot be None"
    # Run performance benchmark
    perf_result = _run_performance_test(
        test,
        implementation_fn,
        context.implementation_path,
        context.reference_path,
        test_input,
        context.use_cached_reference,
        use_defenses=profiling_config.use_defenses,
    )
    profile_traces = []
    ncu_reports = []
    if perf_result.successfully_ran:
        speedup_str = f" ({perf_result.speedup:.2f}x)" if perf_result.speedup else ""
        avg_time_us = perf_result.avg_time_ms * MS_TO_MICROSECONDS
        print(f"   ⏱️  Performance: {avg_time_us:.3f}μs{speedup_str}")
        new_prof_traces, new_ncu_reports = _run_profiling(
            test,
            implementation_fn,
            context,
            profiling_config,
        )
        profile_traces.extend(new_prof_traces)
        ncu_reports.extend(new_ncu_reports)
    else:
        print(f"   ⚠️  Benchmark failed: {perf_result.error_msg}")
    return perf_result, profile_traces, ncu_reports
def _extract_error_type(error_msg: ErrorInfo) -> str:
    """Extract error signature from ErrorInfo.
    Args:
        error_msg: ErrorInfo object
    Returns:
        Error signature for display (e.g., "ValueError: too many values to unpack")
    """
    return error_msg.get_signature()
def _track_error(error_tracker: dict, error_msg: ErrorInfo, test_name: str) -> None:
    """Track error for deduplication.
    Args:
        error_tracker: Dict mapping error signature -> (error_string, test_names[])
        error_msg: ErrorInfo object
        test_name: Name of test that failed
    """
    signature = error_msg.get_signature()
    error_string = error_msg.to_string()
    if signature not in error_tracker:
        error_tracker[signature] = (error_string, [])
    error_tracker[signature][1].append(test_name)
def _print_error_summary(error_tracker: dict) -> None:
    """Print deduplicated error summary.
    Args:
        error_tracker: Dict mapping error signature -> (error_info, test_names[])
    """
    print("\n" + "=" * PRINT_SEPARATOR_WIDTH)
    print("❌ Error Summary")
    print("=" * PRINT_SEPARATOR_WIDTH)
    for idx, (error_type, (error_string, test_names)) in enumerate(error_tracker.items(), 1):
        print(f"\n[Error {idx}] {error_type}")
        print(f"   Affected tests ({len(test_names)}): {', '.join(test_names)}")
        print("\n   Full traceback:")
        for line in error_string.split("\n"):
            print(f"   {line}")
    print("\n" + "=" * PRINT_SEPARATOR_WIDTH)
def _run_single_test(
    test: TestCase,
    test_idx: int,
    total_tests: int,
    implementation_fn: object,
    checker: object | None,
    config: TestRunConfig,
    error_tracker: dict | None = None,
) -> tuple[CorrectnessResult, PerformanceResult, list, list]:
    """Run a single test case with correctness, performance, and profiling.
    Returns:
        (correctness_result, performance_result, profile_traces, ncu_reports)
    """
    assert test is not None, "test cannot be None"
    assert implementation_fn is not None, "implementation_fn cannot be None"
    assert 1 <= test_idx <= total_tests, f"test_idx must be in [1, {total_tests}], got {test_idx}"
    # Tiger Style: Type narrow to Callable after assertion
    assert callable(implementation_fn), "implementation_fn must be callable"
    impl_fn: Callable = implementation_fn
    # Tiger Style: Type narrow checker to Callable | None
    checker_fn: Callable | None = None
    if checker is not None:
        assert callable(checker), "checker must be callable when not None"
        checker_fn = checker
    print(f"\n📊 Test [{test_idx}/{total_tests}]: {test.name}")
    print(f"   Parameters: {test.serialize()}")
    # Type narrowing: checker_fn must not be None (we always use live reference mode)
    assert checker_fn is not None, (
        "checker_fn should not be None when use_cached_reference is False"
    )
    is_correct, error_msg, test_input = _run_correctness_test(
        test, impl_fn, checker_fn, config.context.use_cached_reference, config.context.test_suite
    )
    error_msg_str = error_msg.to_string() if error_msg else None
    correctness_result = CorrectnessResult(
        test_name=test.name,
        backend_name=config.context.implementation_path,
        is_correct=is_correct,
        test_params=test.serialize(),
        error_msg=error_msg_str,
    )
    if is_correct:
        print("   ✅ Correctness: PASS")
        perf_result, profile_traces, ncu_reports = _run_performance_and_profiling(
            test,
            impl_fn,
            test_input,
            config.context,
            config.profiling,
        )
        return correctness_result, perf_result, profile_traces, ncu_reports
    else:
        print("   ❌ Correctness: FAIL")
        if error_msg and error_tracker is not None:
            _track_error(error_tracker, error_msg, test.name)
            # Just print short error type, full traceback will be in summary
            error_type = _extract_error_type(error_msg)
            print(f"   Error: {error_type}")
        elif error_msg:
            # Fallback if no error tracker
            print(f"   Error: {error_msg.to_string()}")
        # Still record performance (as failed)
        perf_result = PerformanceResult(
            test_name=test.name,
            backend_name=config.context.implementation_path,
            successfully_ran=False,
            test_params=test.serialize(),
            avg_time_ms=ERROR_TIME_MS,
            error_msg="Skipped due to correctness failure",
        )
        return correctness_result, perf_result, [], []
def _run_test_loop(
    test_cases: list[TestCase],
    implementation_fn: object,
    checker: object | None,
    config: TestRunConfig,
) -> tuple[list[CorrectnessResult], list[PerformanceResult], list, list, bool]:
    """Run all test cases and collect results.
    Returns:
        (correctness_results, performance_results, profile_traces, ncu_reports, early_stop)
    """
    assert test_cases is not None, "test_cases cannot be None"
    assert len(test_cases) > 0, "test_cases cannot be empty"
    correctness_results = []
    performance_results = []
    profile_traces = []
    ncu_reports = []
    early_stop = False

    error_tracker = {}  # Maps error signature -> (full_error, test_names[])
    for test_idx, test in enumerate(test_cases, 1):
        if early_stop:
            # Skip remaining tests - kernel has fundamental errors
            corr_result, perf_result = _create_skipped_test_results(
                test, config.context.implementation_path
            )
            correctness_results.append(corr_result)
            performance_results.append(perf_result)
            continue
        corr_result, perf_result, new_prof_traces, new_ncu_reports = _run_single_test(
            test,
            test_idx,
            len(test_cases),
            implementation_fn,
            checker,
            config,
            error_tracker,  # Pass error tracker to collect errors
        )
        correctness_results.append(corr_result)
        performance_results.append(perf_result)
        profile_traces.extend(new_prof_traces)
        ncu_reports.extend(new_ncu_reports)
        if not corr_result.is_correct and corr_result.error_msg:
            fundamental_errors = [
                "TypeError",
                "ValueError",
                "AttributeError",
                "NameError",
                "SyntaxError",
                "ImportError",
                "CUDA error",
                "AcceleratorError",
            ]
            # error_msg is now an ErrorInfo object, check exc_type field
            if isinstance(corr_result.error_msg, ErrorInfo):
                error_type = corr_result.error_msg.exc_type
            else:
                error_type = str(corr_result.error_msg)
            if any(err in error_type for err in fundamental_errors):
                early_stop = True
                print("   ⚠️  Stopping early - fundamental error detected")
    # Print error summary at the end (deduplicated)
    if error_tracker:
        _print_error_summary(error_tracker)
    return correctness_results, performance_results, profile_traces, ncu_reports, early_stop
@dataclass(frozen=True)
class ResultsSummary:
    """Summary data for saving and printing results."""
    profile_traces: list
    ncu_reports: list
@dataclass(frozen=True)
class EvaluationArtifactConfig:
    """Configuration for evaluation artifacts and profiling.
    Controls what gets saved and where.
    """
    run_profiling: bool = False
    create_artifact: bool = True
    artifact_name: str | None = None
    run_dir: str | None = None
    use_defenses: bool = False
def _save_and_print_results(
    backend_results: BackendResults,
    test_suites: list[str],
    test_cases: list[TestCase],
    run_path: Path,
    save_results: bool,
    results_summary: ResultsSummary,
    profiling_config: ProfilingConfig,
) -> None:
    """Save results to file and print summaries."""
    assert backend_results is not None, "backend_results cannot be None"
    assert run_path is not None, "run_path cannot be None"
    print(f"\n   {backend_results.summary()}")
    # Print GPUMode format
    print("\n" + "=" * PRINT_SEPARATOR_WIDTH)
    print_gpumode_format(backend_results, test_suites, test_cases)
    # Save results if requested
    if save_results:
        suite_name = "_".join(test_suites).upper()
        suite_results = TestSuiteResults(
            suite_name=suite_name,
            backends=[backend_results],
        )
        output_path = run_path / "results.json"
        suite_results.to_json(output_path)
        print(f"\n💾 Results saved to: {output_path}")
    # Print profile summaries
    if profiling_config.enable_profiling and results_summary.profile_traces:
        print("\n📊 Torch Profiling Summary:")
        print(f"   {len(results_summary.profile_traces)} profile trace(s) generated")
        print("   Location: profiles/")
        print("\n   To view torch profiles:")
        print("   1. Chrome trace: Open chrome://tracing and load .json files")
        print("   2. TensorBoard: tensorboard --logdir=profiles/")
    if profiling_config.enable_ncu and results_summary.ncu_reports:
        print("\n📊 NCU Profiling Summary:")
        print(f"   {len(results_summary.ncu_reports)} NCU report(s) generated")
        print("   Location: ncu_reports/")
        print("\n   To view NCU reports:")
        print("   CSV files in ncu_reports/")
def _print_evaluation_header(
    implementation_path: str,
    reference_path: str,
    test_suites: list[str],
    test_cases: list[TestCase],
    enable_profiling: bool,
    enable_ncu: bool,
) -> None:
    """Print evaluation header information."""
    assert implementation_path, "implementation_path cannot be empty"
    assert test_suites is not None, "test_suites cannot be None"
    assert test_cases is not None, "test_cases cannot be None"
    print("🔥 Running GPU Kernel Tests")
    print(f"   Test Suites: {', '.join(test_suites)}")
    print(f"   Total Tests: {len(test_cases)}")
    print(f"   Implementation: {implementation_path}")
    print(f"   Reference: {reference_path}")
    if enable_profiling:
        print("   Torch Profiling: ENABLED")
    if enable_ncu:
        print("   NCU Profiling: ENABLED")
    print("=" * PRINT_SEPARATOR_WIDTH)
def _execute_test_loop_and_collect_results(
    test_cases: list[TestCase],
    implementation_fn: object,
    checker: object | None,
    config: TestRunConfig,
) -> tuple[BackendResults, list, list]:
    """Execute test loop and create backend results.
    Returns:
        (backend_results, profile_traces, ncu_reports)
    """
    assert test_cases is not None, "test_cases cannot be None"
    assert len(test_cases) > 0, "test_cases cannot be empty"
    assert implementation_fn is not None, "implementation_fn cannot be None"
    correctness_results, performance_results, profile_traces, ncu_reports, _ = _run_test_loop(
        test_cases,
        implementation_fn,
        checker,
        config,
    )
    backend_results = BackendResults(
        backend_name=config.context.implementation_path,
        correctness_tests=correctness_results,
        performance_tests=performance_results,
    )
    return backend_results, profile_traces, ncu_reports
def _setup_and_print_backend_info(
    implementation_path: str,
    reference_path: str,
    test_suites: list[str],
) -> tuple[object, object | None, bool] | tuple[None, None, None]:
    """Setup implementation and print information.
    Returns:
        (implementation_fn, checker, use_cached_reference) on success, (None, None, None) on failure
    """
    assert implementation_path, "implementation_path cannot be empty"
    assert reference_path, "reference_path cannot be empty"
    assert test_suites is not None, "test_suites cannot be None"
    implementation_fn, checker, use_cached_reference, error = _setup_test_backend(
        implementation_path, reference_path, test_suites
    )
    if error:
        print(f"❌ {error}")
        return None, None, None
    print(f"\n🔧 Implementation: {implementation_path}")
    print("-" * PRINT_SEPARATOR_WIDTH)
    return implementation_fn, checker, use_cached_reference
def _validate_evaluation_inputs(
    implementation_path: str,
    reference_path: str,
    test_cases: list[TestCase],
) -> str | None:
    """Validate evaluation inputs, return error message if invalid.
    Tiger Style: Pure validation function, returns error or None.
    Args:
        implementation_path: Path to implementation kernel file
        reference_path: Path to reference kernel file
        test_cases: List of TestCase objects to run
    Returns:
        Error message if validation fails, None if all inputs valid
    """
    assert implementation_path, "implementation_path cannot be empty"
    assert reference_path, "reference_path cannot be empty"
    assert test_cases is not None, "test_cases cannot be None"
    assert len(test_cases) > 0, "test_cases cannot be empty"
    if not torch.cuda.is_available():
        return "CUDA not available - tests require GPU"
    if not Path(implementation_path).exists():
        return f"Implementation file not found: {implementation_path}"
    if not Path(reference_path).exists():
        return f"Reference file not found: {reference_path}"
    return None
def _setup_evaluation_directories(
    run_dir: str | None,
) -> tuple[Path, str, str]:
    """Create run directory structure for evaluation outputs.
    Tiger Style: Pure computation, no branching logic.
    Args:
        run_dir: Directory for this run's outputs. If None, uses run_TIMESTAMP
    Returns:
        (run_path, profile_subdir, ncu_subdir)
    """
    from datetime import datetime

    if run_dir is None:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        actual_run_dir = f"run_{timestamp}"
    else:
        actual_run_dir = run_dir
    run_path = Path(actual_run_dir)
    run_path.mkdir(parents=True, exist_ok=True)
    artifact_path = run_path / "artifact"
    artifact_path.mkdir(parents=True, exist_ok=True)
    (artifact_path / "profiles").mkdir(exist_ok=True)
    (artifact_path / "ncu").mkdir(exist_ok=True)
    (artifact_path / "results").mkdir(exist_ok=True)
    profile_subdir = str(artifact_path / "profiles")
    ncu_subdir = str(artifact_path / "ncu")
    return run_path, profile_subdir, ncu_subdir
def _setup_implementation_backend(
    implementation_path: str,
    reference_path: str,
    test_cases: list[TestCase],
    run_profiling: bool,
) -> tuple[Callable | None, Callable | None, bool]:
    """Setup implementation backend and print evaluation header.
    Tiger Style: Wrapper that handles setup + printing side effects.
    Args:
        implementation_path: Path to implementation kernel
        reference_path: Path to reference kernel
        test_cases: Test cases to display in header
        run_profiling: Whether profiling is enabled (for header display)
    Returns:
        (implementation_fn, checker, use_cached_reference)
        Returns (None, None, False) if setup fails
    """
    test_suites_display = ["custom"]  # For display only
    _print_evaluation_header(
        implementation_path,
        reference_path,
        test_suites_display,
        test_cases,
        run_profiling,
        run_profiling,
    )

    implementation_setup = _setup_and_print_backend_info(
        implementation_path, reference_path, test_suites_display
    )
    if implementation_setup[0] is None:
        return None, None, False
    # Tiger Style: Type narrow - extract and validate components
    impl_fn_obj, checker_obj, use_cached_obj = implementation_setup
    assert callable(impl_fn_obj), "implementation_fn must be callable"
    impl_fn: Callable = impl_fn_obj
    checker_fn: Callable | None = None
    if checker_obj is not None:
        assert callable(checker_obj), "checker must be callable when not None"
        checker_fn = checker_obj
    # Tiger Style: Type narrow use_cached from bool | None to bool
    assert isinstance(use_cached_obj, bool), "use_cached must be bool when impl_fn is not None"
    use_cached: bool = use_cached_obj
    return impl_fn, checker_fn, use_cached
@dataclass
class EvaluationResult:
    """Full evaluation result with structured test data.
    Provides detailed per-test information for better error reporting.
    """
    all_passed: bool
    summary: str
    artifact_path: Path | None
    backend_results: BackendResults | None = None
    error_message: str | None = None
    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        result = {
            "compiled": self.backend_results is not None,
            "all_correct": self.all_passed,
            "error_message": self.error_message,
        }
        if self.backend_results:
            result.update({
                "correctness_score": self.backend_results.correctness_score,
                "geomean_speedup": self.backend_results.geomean_speedup,
                "passed_tests": sum(
                    1 for t in self.backend_results.correctness_tests if t.is_correct
                ),
                "total_tests": len(self.backend_results.correctness_tests),
                "correctness_tests": [t.to_dict() for t in self.backend_results.correctness_tests],
                "performance_tests": [t.to_dict() for t in self.backend_results.performance_tests],
            })
        else:
            result.update({
                "correctness_score": 0.0,
                "geomean_speedup": 0.0,
                "passed_tests": 0,
                "total_tests": 0,
                "correctness_tests": [],
                "performance_tests": [],
            })
        return result
def run_evaluation_full(
    implementation_path: str,
    reference_path: str,
    test_cases: list[TestCase],
    run_benchmarks: bool = False,
    artifact_config: EvaluationArtifactConfig | None = None,
) -> EvaluationResult:
    """Run evaluation and return full structured results.
    Like run_evaluation but returns EvaluationResult with full BackendResults
    for detailed error reporting.
    Args:
        implementation_path: Path to implementation kernel file
        reference_path: Path to reference kernel file
        test_cases: List of TestCase objects to run
        run_benchmarks: Run performance benchmarks (default: False, correctness only)
        artifact_config: Configuration for artifacts, profiling, and output dirs
    Returns:
        EvaluationResult with full test details
    """
    if artifact_config is None:
        artifact_config = EvaluationArtifactConfig()
    validation_error = _validate_evaluation_inputs(implementation_path, reference_path, test_cases)
    if validation_error:
        return EvaluationResult(
            all_passed=False,
            summary=validation_error,
            artifact_path=None,
            backend_results=None,
            error_message=validation_error,
        )

    run_path, profile_subdir, ncu_subdir = _setup_evaluation_directories(artifact_config.run_dir)

    implementation_fn, checker, use_cached_reference = _setup_implementation_backend(
        implementation_path, reference_path, test_cases, artifact_config.run_profiling
    )
    artifact_path = None
    artifact_path_inner = run_path / "artifact" if artifact_config.create_artifact else None
    if implementation_fn is None:
        # Compilation/import failed
        error_msg = "Failed to setup implementation (compilation or import error)"
        if artifact_config.create_artifact:
            assert artifact_path_inner is not None
            artifact_path = _create_submission_artifact(
                implementation_path=implementation_path,
                reference_path=reference_path,
                test_cases=test_cases,
                artifact_path=artifact_path_inner,
                run_dir=str(run_path),
            )
        return EvaluationResult(
            all_passed=False,
            summary=error_msg,
            artifact_path=artifact_path,
            backend_results=None,
            error_message=error_msg,
        )
    if artifact_config.create_artifact:
        assert artifact_path_inner is not None
        artifact_path = _create_submission_artifact_base(
            implementation_path=implementation_path,
            reference_path=reference_path,
            test_cases=test_cases,
            artifact_path=artifact_path_inner,
        )
    # Run tests
    try:
        test_suite_name = "custom"
        context = TestExecutionContext(
            implementation_path=implementation_path,
            reference_path=reference_path,
            test_suite=test_suite_name,
            use_cached_reference=use_cached_reference,
        )
        profiling_config = ProfilingConfig(
            profile_subdir=profile_subdir,
            ncu_subdir=ncu_subdir,
            enable_profiling=artifact_config.run_profiling,
            enable_ncu=artifact_config.run_profiling,
            use_defenses=artifact_config.use_defenses,
        )
        config = TestRunConfig(context=context, profiling=profiling_config)
        backend_results, profile_traces, ncu_reports = _execute_test_loop_and_collect_results(
            test_cases,
            implementation_fn,
            checker,
            config,
        )
        # Save and print results
        test_suites_display = ["custom"]
        results_summary = ResultsSummary(
            profile_traces=profile_traces,
            ncu_reports=ncu_reports,
        )
        _save_and_print_results(
            backend_results,
            test_suites_display,
            test_cases,
            run_path,
            True,
            results_summary,
            profiling_config,
        )
        all_passed = backend_results.all_correct
        summary = f"Implementation '{implementation_path}': {backend_results.summary()}"
        error_message = None
        if not all_passed:
            failed_tests = [t for t in backend_results.correctness_tests if not t.is_correct]
            if failed_tests:
                error_message = failed_tests[0].error_msg  # First failure's error
        return EvaluationResult(
            all_passed=all_passed,
            summary=summary,
            artifact_path=artifact_path,
            backend_results=backend_results,
            error_message=error_message,
        )
    finally:
        if artifact_config.create_artifact and artifact_path:
            _update_artifact_with_results(artifact_path, run_path)
def run_evaluation(
    implementation_path: str,
    reference_path: str,
    test_cases: list[TestCase],
    run_benchmarks: bool = False,
    artifact_config: EvaluationArtifactConfig | None = None,
) -> tuple[bool, str, Path | None]:
    """Run evaluation with explicit test cases and artifact creation.
    New simplified API - accepts test cases directly, no test suite lookups.
    Tiger Style: Push ifs up - parent has control flow, helpers do work.
    Args:
        implementation_path: Path to implementation kernel file
        reference_path: Path to reference kernel file
        test_cases: List of TestCase objects to run
        run_benchmarks: Run performance benchmarks (default: False, correctness only)
        artifact_config: Configuration for artifacts, profiling, and output dirs
    Returns:
        (all_passed, summary_msg, artifact_path)
    Example:
        # Simple usage (default artifact config)
        run_evaluation(impl_path, ref_path, test_cases)
        # With custom artifact config
        cfg = EvaluationArtifactConfig(run_profiling=True, create_artifact=False)
        run_evaluation(impl_path, ref_path, test_cases, artifact_config=cfg)
    """
    if artifact_config is None:
        artifact_config = EvaluationArtifactConfig()
    validation_error = _validate_evaluation_inputs(implementation_path, reference_path, test_cases)
    if validation_error:
        return False, validation_error, None

    run_path, profile_subdir, ncu_subdir = _setup_evaluation_directories(artifact_config.run_dir)

    implementation_fn, checker, use_cached_reference = _setup_implementation_backend(
        implementation_path, reference_path, test_cases, artifact_config.run_profiling
    )
    artifact_path = None
    artifact_path_inner = run_path / "artifact" if artifact_config.create_artifact else None
    if implementation_fn is None:
        # Compilation/import failed - create artifact with error info
        if artifact_config.create_artifact:
            # Tiger Style: Assert invariant - artifact_path_inner set when create_artifact=True
            assert artifact_path_inner is not None, (
                "artifact_path_inner must be set when create_artifact=True"
            )
            artifact_path = _create_submission_artifact(
                implementation_path=implementation_path,
                reference_path=reference_path,
                test_cases=test_cases,
                artifact_path=artifact_path_inner,
                run_dir=str(run_path),
            )
        return False, "Failed to setup implementation", artifact_path
    # CRITICAL: Create artifact BEFORE running tests
    # This ensures we capture source files even if tests crash
    if artifact_config.create_artifact:
        # Tiger Style: Assert invariant - artifact_path_inner set when create_artifact=True
        assert artifact_path_inner is not None, (
            "artifact_path_inner must be set when create_artifact=True"
        )
        artifact_path = _create_submission_artifact_base(
            implementation_path=implementation_path,
            reference_path=reference_path,
            test_cases=test_cases,
            artifact_path=artifact_path_inner,
        )
    # Run tests with try/finally to ensure results are saved even on crash
    try:
        test_suite_name = "custom"  # For naming only
        context = TestExecutionContext(
            implementation_path=implementation_path,
            reference_path=reference_path,
            test_suite=test_suite_name,
            use_cached_reference=use_cached_reference,
        )
        profiling_config = ProfilingConfig(
            profile_subdir=profile_subdir,
            ncu_subdir=ncu_subdir,
            enable_profiling=artifact_config.run_profiling,
            enable_ncu=artifact_config.run_profiling,
            use_defenses=artifact_config.use_defenses,
        )
        config = TestRunConfig(context=context, profiling=profiling_config)
        # Run tests
        backend_results, profile_traces, ncu_reports = _execute_test_loop_and_collect_results(
            test_cases,
            implementation_fn,
            checker,
            config,
        )
        # Save and print results
        test_suites_display = ["custom"]
        results_summary = ResultsSummary(
            profile_traces=profile_traces,
            ncu_reports=ncu_reports,
        )
        _save_and_print_results(
            backend_results,
            test_suites_display,
            test_cases,
            run_path,
            True,  # Always save results
            results_summary,
            profiling_config,
        )
        all_passed = backend_results.all_correct
        summary = f"Implementation '{implementation_path}': {backend_results.summary()}"
        return all_passed, summary, artifact_path
    finally:
        # ALWAYS update artifact with results and logs, even on crash
        if artifact_config.create_artifact and artifact_path:
            _update_artifact_with_results(artifact_path, run_path)
def _capture_environment_info(artifact_path: Path) -> None:
    """Capture Python and GPU environment information.
    Tries nvidia-smi first, falls back to PyTorch/CUDA info if nvidia-smi fails.
    This handles cases where NVML has driver/library version mismatch but CUDA still works.
    Tiger Style: Graceful degradation with fallback.
    Args:
        artifact_path: Artifact directory to save environment info
    """
    import subprocess
    import sys
    # Save Python version
    (artifact_path / "environment" / "python_version.txt").write_text(sys.version)
    # Try to capture GPU info via nvidia-smi or rocm-smi
    gpu_smi_output = None
    gpu_smi_file = "gpu_smi.txt"
    # Try nvidia-smi first (for NVIDIA GPUs)
    try:
        nvidia_smi_result = subprocess.run(
            ["nvidia-smi"],
            capture_output=True,
            text=True,
            timeout=10,
            check=False,
        )
        if nvidia_smi_result.returncode == 0:
            gpu_smi_output = nvidia_smi_result.stdout
            gpu_smi_file = "nvidia_smi.txt"
    except FileNotFoundError:
        pass  # nvidia-smi not available, try rocm-smi
    # Try rocm-smi if nvidia-smi didn't work (for AMD GPUs)
    if gpu_smi_output is None:
        try:
            rocm_smi_result = subprocess.run(
                ["rocm-smi"],
                capture_output=True,
                text=True,
                timeout=10,
                check=False,
            )
            if rocm_smi_result.returncode == 0:
                gpu_smi_output = rocm_smi_result.stdout
                gpu_smi_file = "rocm_smi.txt"
        except FileNotFoundError:
            pass  # rocm-smi not available either
    if gpu_smi_output:
        (artifact_path / "environment" / gpu_smi_file).write_text(gpu_smi_output)
    else:
        # Neither nvidia-smi nor rocm-smi available - fall back to PyTorch info
        (artifact_path / "environment" / "gpu_smi.txt").write_text(
            "Neither nvidia-smi nor rocm-smi available.\n"
            "Falling back to PyTorch GPU info..."
        )
        # Collect GPU info via PyTorch as fallback
        try:
            import torch
            gpu_info = []
            gpu_info.append(f"PyTorch version: {torch.__version__}")
            gpu_info.append(f"CUDA available: {torch.cuda.is_available()}")
            if torch.cuda.is_available():
                gpu_info.append(f"CUDA version: {torch.version.cuda}")
                gpu_info.append(f"cuDNN version: {torch.backends.cudnn.version()}")
                gpu_info.append(f"Device count: {torch.cuda.device_count()}")
                for i in range(torch.cuda.device_count()):
                    props = torch.cuda.get_device_properties(i)
                    gpu_info.append(f"\nGPU {i}: {props.name}")
                    gpu_info.append(f"  Compute capability: {props.major}.{props.minor}")
                    gpu_info.append(f"  Total memory: {props.total_memory / 1024**3:.2f} GB")
                    gpu_info.append(f"  Multi-processor count: {props.multi_processor_count}")
            fallback_info = "\n".join(gpu_info)
            (artifact_path / "environment" / "gpu_info_pytorch.txt").write_text(fallback_info)
        except Exception as e:
            # If even PyTorch fails, save the error
            (artifact_path / "environment" / "gpu_info_pytorch.txt").write_text(
                f"Failed to collect GPU info via PyTorch: {e}"
            )
def _copy_results_to_artifact(artifact_path: Path, run_dir: str | None) -> None:
    """Copy evaluation results to artifact directory.
    Tiger Style: Pure side-effect function.
    Args:
        artifact_path: Artifact directory
        run_dir: Run directory containing results.json
    """
    if not run_dir:
        return
    run_path = Path(run_dir)
    if not run_path.exists():
        return
    results_file = run_path / "results.json"
    if results_file.exists():
        import shutil
        shutil.copy2(results_file, artifact_path / "results" / "full_results.json")
def _generate_artifact_manifest(artifact_path: Path, implementation_path: str) -> None:
    """Generate manifest.json with artifact metadata.
    Tiger Style: Pure side-effect function.
    Args:
        artifact_path: Artifact directory
        implementation_path: Path to implementation (for naming)
    """
    import json
    from datetime import datetime, timezone
    manifest = {
        "submission_name": Path(implementation_path).stem,
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "files": {
            "submission": "submission.py",
            "reference": "reference.py",
            "test_cases": "tests.json",
        },
    }
    with open(artifact_path / "manifest.json", "w") as f:
        json.dump(manifest, f, indent=2)
def _create_submission_artifact_base(
    implementation_path: str,
    reference_path: str,
    test_cases: list[TestCase],
    artifact_path: Path,
) -> Path:
    """Create base artifact with source files BEFORE tests run.
    This is called immediately after kernel loading succeeds, ensuring
    we capture all inputs even if tests crash during execution.
    Tiger Style: Early fail-safe artifact creation.
    Args:
        implementation_path: Path to implementation kernel
        reference_path: Path to reference kernel
        test_cases: Test cases to run
        artifact_path: Path to artifact directory (already created)
    Returns:
        Path to artifact directory
    Raises:
        FileNotFoundError: If source files not found
        IOError: If file operations fail
    """
    import json
    import shutil
    (artifact_path / "environment").mkdir(exist_ok=True)
    (artifact_path / "logs").mkdir(exist_ok=True)
    # Copy source files - these are the INPUTS to evaluation
    shutil.copy2(implementation_path, artifact_path / "submission.py")
    shutil.copy2(reference_path, artifact_path / "reference.py")
    # Save test cases
    test_cases_data = [{**tc.params, "name": tc.name} for tc in test_cases]
    with open(artifact_path / "tests.json", "w") as f:
        json.dump(test_cases_data, f, indent=2)
    # Capture environment info
    _capture_environment_info(artifact_path)
    # Generate manifest
    _generate_artifact_manifest(artifact_path, implementation_path)
    return artifact_path
def _update_artifact_with_results(artifact_path: Path, run_dir: Path) -> None:
    """Update artifact with results and logs after tests complete or crash.
    Called in finally block to ensure results are saved even on crash.
    Safe to call multiple times (idempotent).
    Tiger Style: Defensive - handles missing files gracefully.
    Args:
        artifact_path: Path to artifact directory
        run_dir: Run directory with results/logs
    """
    import shutil
    # Copy results.json if it exists
    results_file = run_dir / "results.json"
    if results_file.exists():
        try:
            shutil.copy2(results_file, artifact_path / "results" / "full_results.json")
        except Exception as e:
            # Non-fatal - log but don't fail
            print(f"Warning: Failed to copy results.json: {e}")
    # Copy evaluate.log if it exists (from tmux session)
    log_file = run_dir / "evaluate.log"
    if log_file.exists():
        try:
            shutil.copy2(log_file, artifact_path / "logs" / "evaluate.log")
        except Exception as e:
            # Non-fatal - log but don't fail
            print(f"Warning: Failed to copy evaluate.log: {e}")
def _create_submission_artifact(
    implementation_path: str,
    reference_path: str,
    test_cases: list[TestCase],
    artifact_path: Path,
    run_dir: str | None,
) -> Path:
    """Create submission artifact with code, environment, and results.
    Tiger Style: Orchestrator - sequential helper calls, explicit errors.
    NOTE: This is the OLD function kept for backward compatibility.
    New code should use _create_submission_artifact_base() + _update_artifact_with_results()
    Args:
        implementation_path: Path to implementation kernel
        reference_path: Path to reference kernel
        test_cases: Test cases that were run
        artifact_path: Path to artifact directory (already created)
        run_dir: Run directory with results
    Returns:
        Path to created artifact directory
    Raises:
        FileNotFoundError: If nvidia-smi missing or source files not found
    """
    import json
    import shutil
    (artifact_path / "environment").mkdir(exist_ok=True)
    (artifact_path / "logs").mkdir(exist_ok=True)
    # Copy source files
    shutil.copy2(implementation_path, artifact_path / "submission.py")
    shutil.copy2(reference_path, artifact_path / "reference.py")
    # Save test cases
    test_cases_data = [{**tc.params, "name": tc.name} for tc in test_cases]
    with open(artifact_path / "tests.json", "w") as f:
        json.dump(test_cases_data, f, indent=2)
    # Capture environment (fast-fails if GPU misconfigured)
    _capture_environment_info(artifact_path)
    # Copy results
    _copy_results_to_artifact(artifact_path, run_dir)
    # Generate manifest
    _generate_artifact_manifest(artifact_path, implementation_path)
    return artifact_path
def load_kernel_from_file(path: str, function_name: str = "custom_kernel") -> Callable:
    """Load kernel function from Python file.
    Uses GPU Mode's standard naming convention:
    - Submission files must define 'custom_kernel'
    - Reference files must define 'ref_kernel'
    Args:
        path: Path to Python file containing kernel function
        function_name: Name of function to load (default: "custom_kernel")
    Returns:
        Kernel function object
    Raises:
        AssertionError: If file doesn't exist or doesn't contain the required function
    """
    import importlib.util
    from pathlib import Path
    assert path, "path cannot be empty"
    assert function_name, "function_name cannot be empty"
    kernel_file = Path(path)
    assert kernel_file.exists(), f"Kernel file not found: {path}"
    # This is needed for GPUMode reference kernels that do: from task import input_t, output_t
    kernel_dir = str(kernel_file.parent.resolve())
    sys.path.insert(0, kernel_dir)
    try:
        spec = importlib.util.spec_from_file_location("dynamic_kernel", path)
        assert spec is not None, f"Failed to load spec from: {path}"
        assert spec.loader is not None, f"No loader for spec: {path}"
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        assert hasattr(module, function_name), f"File {path} must define {function_name} function"
        return getattr(module, function_name)
    finally:
        # Remove from sys.path to avoid polluting it
        if kernel_dir in sys.path:
            sys.path.remove(kernel_dir)
def load_test_cases_from_json(test_cases_arg: str) -> list[TestCase]:
    """Load test cases from JSON file or inline JSON string.
    Args:
        test_cases_arg: Path to JSON file or inline JSON string
    Returns:
        List of TestCase objects
    """
    import json
    from pathlib import Path
    assert test_cases_arg, "test_cases_arg cannot be empty"
    # Try to parse as file path first
    test_cases_path = Path(test_cases_arg)
    if test_cases_path.exists() and test_cases_path.is_file():
        with open(test_cases_path) as f:
            test_cases_data = json.load(f)
    else:
        # Try to parse as inline JSON
        try:
            test_cases_data = json.loads(test_cases_arg)
            # If single dict, wrap in list
            if isinstance(test_cases_data, dict):
                test_cases_data = [test_cases_data]
        except json.JSONDecodeError as e:
            raise InvalidTestCasesFormatError(test_cases_arg, e) from e
    # Validate it's a list
    assert isinstance(test_cases_data, list), (
        f"Test cases must be a JSON array, got {type(test_cases_data)}"
    )
    assert len(test_cases_data) > 0, "Test cases cannot be empty"
    # All test cases are now generic - just dict of params
    test_cases = []
    for i, tc_dict in enumerate(test_cases_data):
        assert isinstance(tc_dict, dict), f"Test case {i} must be a dict, got {type(tc_dict)}"
        name = tc_dict.get("name", f"test_{i + 1}")
        # All params except 'name' go into the params dict
        params = {k: v for k, v in tc_dict.items() if k != "name"}
        test_cases.append(
            TestCase(
                params=params,
                name=str(name),
            )
        )
    return test_cases
# Protocol for CLI arguments
# Tiger Style: Explicit type instead of 'object'
class EvaluateArgs(Protocol):
    """Protocol for evaluate.py CLI arguments."""
    implementation: str
    reference: str
    test_cases: str
    benchmark: bool
    profile: bool
    defensive: bool
    no_artifact: bool
    artifact_name: str | None
    run_dir: str | None
def _parse_arguments() -> EvaluateArgs:
    """Parse command-line arguments.
    Returns:
        Parsed arguments (EvaluateArgs protocol)
    """
    import argparse
    parser = argparse.ArgumentParser(
        description="Run GPU kernel evaluation with correctness and performance testing"
    )
    parser.add_argument(
        "--implementation",
        "--impl",
        dest="implementation",
        type=str,
        required=True,
        help="Path to implementation kernel file (e.g., my_kernel.py)",
    )
    parser.add_argument(
        "--reference",
        type=str,
        required=True,
        help="Path to reference kernel file (e.g., nvfp4_gemv_blackwell/reference_kernel.py)",
    )
    parser.add_argument(
        "--test-cases",
        type=str,
        required=True,
        help="Path to test cases JSON file or inline JSON string",
    )
    parser.add_argument(
        "--benchmark",
        action="store_true",
        help="Run performance benchmarks (in addition to correctness)",
    )
    parser.add_argument(
        "--profile",
        action="store_true",
        help="Enable profiling with torch.profiler and NCU",
    )
    parser.add_argument(
        "--defensive",
        action="store_true",
        help="Enable defensive timing to detect evaluation hacking (stream injection, etc.)",
    )
    parser.add_argument(
        "--no-artifact",
        action="store_true",
        help="Disable artifact creation (for quick iteration)",
    )
    parser.add_argument(
        "--artifact-name",
        type=str,
        help="Custom artifact name (default: implementation filename)",
    )
    parser.add_argument(
        "--run-dir",
        type=str,
        help="Directory for this run's outputs (results, profiles, artifacts). Default: run_TIMESTAMP",
    )
    # Cast to Protocol for type safety
    parsed_args: EvaluateArgs = parser.parse_args()  # type: ignore[assignment]
    return parsed_args
def _run_evaluation_with_args(args: EvaluateArgs) -> tuple[bool, str, Path | None]:
    """Run evaluation with parsed arguments.
    Args:
        args: Parsed CLI arguments (EvaluateArgs protocol)
    Returns:
        (success, summary, artifact_path): True if all tests passed, summary string, artifact path
    """
    test_cases = load_test_cases_from_json(args.test_cases)
    artifact_config = EvaluationArtifactConfig(
        run_profiling=args.profile,
        create_artifact=not args.no_artifact,
        artifact_name=args.artifact_name,
        run_dir=args.run_dir,
        use_defenses=args.defensive,
    )
    return run_evaluation(
        implementation_path=args.implementation,
        reference_path=args.reference,
        test_cases=test_cases,
        run_benchmarks=args.benchmark,
        artifact_config=artifact_config,
    )
def main() -> int:
    """Main entry point."""
    args = _parse_arguments()
    # Validate implementation file exists
    impl_path = Path(args.implementation)
    if not impl_path.exists():
        print(f"❌ Implementation file not found: {args.implementation}")
        return 1
    # Validate reference file exists
    ref_path = Path(args.reference)
    if not ref_path.exists():
        print(f"❌ Reference file not found: {args.reference}")
        return 1
    # Run tests
    try:
        success, summary, artifact_path = _run_evaluation_with_args(args)
        if artifact_path:
            print(f"\n📁 Artifact created: {artifact_path}")
        if success:
            print(f"\n✅ All tests passed: {summary}")
            return 0
        else:
            print(f"\n❌ Some tests failed: {summary}")
            return 1
    except Exception as e:
        print(f"\n❌ Fatal error: {e}")
        import traceback
        traceback.print_exc()
        return 1
if __name__ == "__main__":
    sys.exit(main())
